From 1f8997398eab139f95159d7e2dba8ac88b14ff08 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 04:20:06 +0000 Subject: [PATCH] refactor(rust): extract the host coroutine into its own crate (#43129) Move the generic async coroutine out of the host crate into litellm-coroutine, with its requirements documented in the crate's AGENTS.md. RouteMachine becomes CallMachine on top of it, host protocol ops move to host/src/protocol.rs, and the messages, OCR, host-python, python-bridge and legacy callbacks crates adopt the new types. Adds error definition rules to litellm-rust/AGENTS.md. Co-authored-by: Yujong Lee --- litellm-rust/AGENTS.md | 8 + litellm-rust/Cargo.lock | 11 + litellm-rust/Cargo.toml | 1 + .../callbacks-legacy-python/src/call.rs | 14 +- .../crates/core/src/messages/route.rs | 49 +-- litellm-rust/crates/core/src/ocr/document.rs | 6 +- litellm-rust/crates/core/src/ocr/route.rs | 343 +++++---------- litellm-rust/crates/core/src/ocr/types.rs | 9 - litellm-rust/crates/coroutine/AGENTS.md | 31 ++ litellm-rust/crates/coroutine/Cargo.toml | 15 + litellm-rust/crates/coroutine/src/co.rs | 42 ++ .../crates/coroutine/src/coroutine.rs | 94 ++++ litellm-rust/crates/coroutine/src/error.rs | 14 + litellm-rust/crates/coroutine/src/lib.rs | 12 + litellm-rust/crates/coroutine/src/reply.rs | 60 +++ .../crates/coroutine/tests/coroutine.rs | 256 +++++++++++ litellm-rust/crates/host-python/AGENTS.md | 4 +- litellm-rust/crates/host-python/Cargo.toml | 1 + .../crates/host-python/src/adapter.rs | 39 +- litellm-rust/crates/host-python/src/driver.rs | 401 +++++++++--------- .../crates/host-python/src/file_reader.rs | 241 +++++++++++ litellm-rust/crates/host-python/src/lib.rs | 6 +- litellm-rust/crates/host/Cargo.toml | 1 + litellm-rust/crates/host/src/host.rs | 38 +- litellm-rust/crates/host/src/lib.rs | 7 +- litellm-rust/crates/host/src/machine/auth.rs | 27 +- .../crates/host/src/machine/call_machine.rs | 137 ++++++ litellm-rust/crates/host/src/machine/mod.rs | 32 +- .../crates/host/src/machine/route_machine.rs | 199 --------- litellm-rust/crates/host/src/protocol.rs | 17 + litellm-rust/crates/host/src/route.rs | 14 - litellm-rust/crates/host/src/run.rs | 147 ++++--- .../crates/llms/src/base_llm/ocr/error.rs | 1 - .../python-bridge/src/logger/machine.rs | 11 +- .../crates/python-bridge/src/logger/tests.rs | 13 +- .../python-bridge/src/routes/messages/host.rs | 33 +- .../python-bridge/src/routes/messages/mod.rs | 4 +- .../python-bridge/src/routes/ocr/document.rs | 247 +++-------- .../python-bridge/src/routes/ocr/host.rs | 100 ++--- .../python-bridge/src/routes/ocr/mod.rs | 4 +- .../python-bridge/src/routes/ocr/project.rs | 110 +++-- 41 files changed, 1660 insertions(+), 1139 deletions(-) create mode 100644 litellm-rust/crates/coroutine/AGENTS.md create mode 100644 litellm-rust/crates/coroutine/Cargo.toml create mode 100644 litellm-rust/crates/coroutine/src/co.rs create mode 100644 litellm-rust/crates/coroutine/src/coroutine.rs create mode 100644 litellm-rust/crates/coroutine/src/error.rs create mode 100644 litellm-rust/crates/coroutine/src/lib.rs create mode 100644 litellm-rust/crates/coroutine/src/reply.rs create mode 100644 litellm-rust/crates/coroutine/tests/coroutine.rs create mode 100644 litellm-rust/crates/host-python/src/file_reader.rs create mode 100644 litellm-rust/crates/host/src/machine/call_machine.rs delete mode 100644 litellm-rust/crates/host/src/machine/route_machine.rs create mode 100644 litellm-rust/crates/host/src/protocol.rs delete mode 100644 litellm-rust/crates/host/src/route.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e5ffcd1c57a..70fcc367905 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -8,3 +8,11 @@ - Split a mixed test file along that line instead of widening visibility to move it - A test for another crate's item belongs in that crate, not in a downstream one - Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own + +## Error definitions + +- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs` +- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each +- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string +- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return +- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0a91f0759c2..c522bf205b4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3013,6 +3013,15 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-coroutine" +version = "0.1.0" +dependencies = [ + "rstest", + "thiserror 2.0.19", + "tokio", +] + [[package]] name = "litellm-cost" version = "0.1.0" @@ -3040,6 +3049,7 @@ name = "litellm-host" version = "0.1.0" dependencies = [ "litellm-auth", + "litellm-coroutine", "rstest", "serde_json", "tokio", @@ -3049,6 +3059,7 @@ dependencies = [ name = "litellm-host-python" version = "0.1.0" dependencies = [ + "bytes", "futures-util", "litellm-host", "pyo3", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 9d05c8d2b98..0c7236e807e 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm" litellm-tracing = { path = "crates/tracing" } tracing = "0.1" litellm-core = { path = "crates/core" } +litellm-coroutine = { path = "crates/coroutine" } litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framing = { path = "crates/framer" } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/call.rs b/litellm-rust/crates/callbacks-legacy-python/src/call.rs index b37790f60a8..9b921070839 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/call.rs @@ -3,8 +3,8 @@ //! lifetime. No other callback host has that obligation, which is why nothing outside //! this crate holds them. -use litellm_host::{machine::Machine, route::Route}; -use litellm_host_python::{RouteHost, lookup, run_call}; +use litellm_host::{machine::Machine, protocol::Protocol}; +use litellm_host_python::{ProtocolHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -63,25 +63,25 @@ impl PublicCall { } } -/// Runs one native call under the legacy `Logging` contract: the route host projects from +/// Runs one native call under the legacy `Logging` contract: the protocol host projects from /// the keyword view the contract prepares, and the contract observes the call. pub fn run_legacy_call( py: Python<'_>, surface: LegacySurface, call: PublicCall, machine: M, - route: H, + host: H, asynchronous: bool, ) -> PyResult> where - H: RouteHost + 'static, - M: Machine::Response> + 'static, + H: ProtocolHost + 'static, + M: Machine::Response> + 'static, { let arguments = call.kwargs.clone_ref(py); run_call( py, machine, - route, + host, Box::new(LegacyLogging::new(py, surface, call, asynchronous)), arguments, asynchronous, diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 8cd3eaf3aa3..fc1a9b63252 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -1,4 +1,5 @@ use std::{ + convert::Infallible, sync::{Arc, Mutex}, time::Duration, }; @@ -9,8 +10,8 @@ use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, - machine::{HostChannel, MachineFault, RouteMachine}, - route::Route, + machine::{CallMachine, HostChannel, MachineFault}, + protocol::Protocol, }; use litellm_secrets::source::SecretSource; use litellm_types::{ @@ -28,15 +29,6 @@ use super::{ }; use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum MessagesOp { - ProjectRequest, -} - -pub enum MessagesOpResult { - Request(Box), -} - /// The caller's request as the host projects it. pub struct MessagesCall { pub model: String, @@ -64,11 +56,11 @@ pub enum MessagesOutput { pub struct Messages; -impl Route for Messages { +impl Protocol for Messages { type Response = MessagesOutput; type Error = Error; - type Op = MessagesOp; - type OpResult = MessagesOpResult; + type Projection = MessagesCall; + type Op = Infallible; type Chunk = Bytes; type StreamHead = (); } @@ -78,13 +70,12 @@ impl From for Error { Self::InvalidRequest(match fault { MachineFault::Abandoned => "messages host driver was abandoned".into(), MachineFault::Protocol(message) => format!("messages {message}"), - MachineFault::Mismatch => "invalid messages host operation result".into(), }) } } pub type MessagesHost = HostChannel; -pub type MessagesMachine = RouteMachine; +pub type MessagesMachine = CallMachine; /// Whether this route serves the request, decided before any callback runs so a host /// can still run its own path. @@ -114,30 +105,28 @@ impl LocalMessagesHost { } impl Host for LocalMessagesHost { - async fn route(&self, op: MessagesOp) -> Result { - match op { - MessagesOp::ProjectRequest => self - .call - .lock() - .unwrap_or_else(|error| error.into_inner()) - .take() - .map(|call| MessagesOpResult::Request(Box::new(call))) - .ok_or_else(|| { - Error::InvalidRequest("messages request was already projected".into()) - }), - } + async fn project(&self) -> Result { + self.call + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .ok_or_else(|| Error::InvalidRequest("messages request was already projected".into())) + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} } } pub fn messages_machine(secrets: Arc) -> MessagesMachine { - RouteMachine::new(move |host| Box::pin(execute(host, secrets.clone()))) + CallMachine::new(move |host| Box::pin(execute(host, secrets.clone()))) } async fn execute( host: MessagesHost, secrets: Arc, ) -> Result { - let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?; + let call = host.project().await?; let stream = call.streams(); let resolved = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?; let secrets = secrets.resolve(resolved.config.secret_names()).await?; diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index ffa4f045e8e..b78c09298de 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -23,9 +23,6 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result { file_name.as_deref(), mime_type.as_deref(), )?), - OcrDocumentInput::HostReader { .. } => Err(Error::InvalidRequest( - "OCR file reader was not read by the host".into(), - )), } } @@ -207,7 +204,7 @@ mod tests { } #[test] - fn byte_documents_are_encoded_and_host_readers_must_be_read_first() { + fn byte_documents_are_encoded() { assert_eq!( prepare_document(OcrDocumentInput::Bytes { bytes: b"abc".as_slice().into(), @@ -217,7 +214,6 @@ mod tests { .unwrap(), document("data:application/pdf;base64,YWJj") ); - assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err()); } #[test] diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index 7f83291bdab..e3e57bd2d77 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -3,102 +3,73 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, - machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, - route::Route, + host::Reply, + machine::{CallMachine, HostChannel, HostTokenProvider, TokenProtocol}, + protocol::Protocol, }; use litellm_llms::base_llm::ocr::{ error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse, }; use super::handler::perform_ocr_request; -use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}; +use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrOp { - ProjectRequest, - ReadDocument, - AcquireAzureAdToken, + AcquireAzureAdToken(Reply), } -pub enum OcrOpResult { - Request { - request: Box>, - caller_token: bool, - }, - Document(OcrFileContent), - AzureAdToken(ResolvedCredential), +/// The caller's request as the host projects it. +pub struct OcrProjection { + pub request: LiteLLMOcrRequest, + /// The caller passed its own Azure AD token provider, which the host keeps. + pub caller_token: bool, } pub struct Ocr; -impl Route for Ocr { +impl Protocol for Ocr { type Response = LiteLLMOcrResponse; type Error = Error; + type Projection = OcrProjection; type Op = OcrOp; - type OpResult = OcrOpResult; type Chunk = std::convert::Infallible; type StreamHead = std::convert::Infallible; } -impl TokenRoute for Ocr { - fn acquire_token_op() -> OcrOp { - OcrOp::AcquireAzureAdToken - } - - fn token_credential(result: OcrOpResult) -> Option { - match result { - OcrOpResult::AzureAdToken(credential) => Some(credential), - _ => None, - } +impl TokenProtocol for Ocr { + fn acquire_token_op(reply: Reply) -> OcrOp { + OcrOp::AcquireAzureAdToken(reply) } } pub type OcrHost = HostChannel; -pub type OcrMachine = RouteMachine; +pub type OcrMachine = CallMachine; -/// The OCR call as a machine: projection, document reading and token acquisition are -/// host operations; everything else runs in Rust. +/// The OCR call as a machine: projection and token acquisition are host operations; +/// everything else runs in Rust. pub fn ocr_machine(client: OcrClient) -> OcrMachine { - RouteMachine::new(move |host| Box::pin(execute(client, host))) + CallMachine::new(move |host| Box::pin(execute(client, host))) } async fn execute(client: OcrClient, host: OcrHost) -> Result { - let OcrOpResult::Request { + let OcrProjection { request, caller_token, - } = host.route(OcrOp::ProjectRequest).await? - else { - return Err(MachineFault::Mismatch.into()); - }; + } = host.project().await?; let request = LiteLLMOcrRequest { azure_ad_token_provider: caller_token .then(|| HostTokenProvider::handle(host.clone())) .or(request.azure_ad_token_provider), - ..*request + ..request }; let caller_document = matches!(request.document, OcrDocumentInput::Document(_)); - let request = prepare_request_document(request, &host).await?; + let request = prepare_request_document(request).await?; perform_ocr_request(&client, request, &host, caller_document).await } async fn prepare_request_document( request: LiteLLMOcrRequest, - host: &OcrHost, ) -> Result { - let request = match &request.document { - OcrDocumentInput::HostReader { mime_type } => { - let mime_type = mime_type.clone(); - let OcrOpResult::Document(content) = host.route(OcrOp::ReadDocument).await? else { - return Err(MachineFault::Mismatch.into()); - }; - request.with_document(OcrDocumentInput::Bytes { - bytes: content.bytes, - file_name: content.file_name, - mime_type, - }) - } - _ => request, - }; if let OcrDocumentInput::Document(_) = &request.document { return request.map_document(super::document::prepare_document); } @@ -107,7 +78,6 @@ async fn prepare_request_document( .map_err(|error| Error::DocumentTask(Arc::new(error)))? } -type Reader = Box Result + Send + Sync>; type BeforeSend = Box Result + Send + Sync>; type Observer = Box; @@ -116,7 +86,6 @@ type Observer = Box; /// projection, and the optional observer sees and may rewrite the wire request. pub struct LocalOcrHost { request: Mutex>>, - reader: Option, before_send: Option, observer: Option, } @@ -125,22 +94,11 @@ impl LocalOcrHost { pub fn new(request: LiteLLMOcrRequest) -> Self { Self { request: Mutex::new(Some(request)), - reader: None, before_send: None, observer: None, } } - pub fn with_reader( - self, - reader: impl Fn() -> Result + Send + Sync + 'static, - ) -> Self { - Self { - reader: Some(Box::new(reader)), - ..self - } - } - pub fn with_before_send( self, before_send: impl Fn(WireRequest, &RequestContext) -> Result @@ -163,25 +121,21 @@ impl LocalOcrHost { } impl litellm_host::host::Host for LocalOcrHost { - async fn route(&self, op: OcrOp) -> Result { + async fn project(&self) -> Result { + self.request + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .map(|request| OcrProjection { + request, + caller_token: false, + }) + .ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into())) + } + + async fn custom_op(&self, op: OcrOp) -> Result<(), Error> { match op { - OcrOp::ProjectRequest => self - .request - .lock() - .unwrap_or_else(|error| error.into_inner()) - .take() - .map(|request| OcrOpResult::Request { - request: Box::new(request), - caller_token: false, - }) - .ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into())), - OcrOp::ReadDocument => self - .reader - .as_ref() - .ok_or_else(|| Error::InvalidRequest("OCR host has no document reader".into())) - .and_then(|reader| reader()) - .map(OcrOpResult::Document), - OcrOp::AcquireAzureAdToken => { + OcrOp::AcquireAzureAdToken(_) => { Err(Error::Auth(litellm_auth::Error::AzureTokenAcquisition( "OCR host has no Azure AD token provider".into(), ))) @@ -2757,7 +2711,7 @@ pub(crate) mod tests { use litellm_auth_gcp::VertexAuth; use litellm_host::{ event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp, HostResult}, + host::{Host, HostOp}, machine::{HostFailure, Machine, MachineStep}, }; use litellm_http::{ @@ -2776,7 +2730,7 @@ pub(crate) mod tests { use rstest::rstest; use serde_json::{Value, json}; - use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; + use crate::ocr::route::{LocalOcrHost, OcrOp, OcrProjection, ocr_machine}; use crate::ocr::{ test_support::{ MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, @@ -3212,42 +3166,42 @@ pub(crate) mod tests { crate::ocr::route::OcrMachine, ) { let mut machine = ocr_machine(client); - let mut result = None; let mut ops = Vec::new(); let outcome = loop { - let op = match machine.resume(result.take()).await { + let op = match machine.resume().await { Ok(MachineStep::Host(op)) => op, Ok(MachineStep::Complete(response)) => break Ok(response), Err(error) => break Err(error), }; let answer = match op { - HostOp::Route(op) => { - ops.push(match op { - OcrOp::ProjectRequest => "ProjectRequest", - OcrOp::ReadDocument => "ReadDocument", - OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", - }); - host.route(op) + HostOp::Project(reply) => { + ops.push("Project"); + host.project() .await - .map(HostResult::Route) + .map(|projection| reply.send(projection)) .map_err(HostFailure::Error) } - HostOp::BeforeSend { wire, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) + HostOp::Custom(op) => { + ops.push(match op { + OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken", + }); + host.custom_op(op).await.map_err(HostFailure::Error) } - HostOp::Emit(event) => { + HostOp::BeforeSend { wire, reply, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| reply.send(wire)) + } + HostOp::Emit(event, reply) => { let event = CallEvent::Machine(event); ops.push(event_name(&event)); host.emit(&event) .await - .map(|()| HostResult::Emitted) + .map(|()| reply.send(())) .map_err(HostFailure::Error) } }; - match answer { - Ok(answer) => result = Some(answer), - Err(failure) => break machine.interrupt(failure).await, + if let Err(failure) = answer { + break machine.interrupt(failure).await; } }; (outcome, ops, machine) @@ -3269,8 +3223,8 @@ pub(crate) mod tests { assert!( matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") ); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(None).await.is_err()); + assert_eq!(ops, ["Project", "BeforeSend"]); + assert!(machine.resume().await.is_err()); } #[tokio::test] @@ -3306,80 +3260,24 @@ pub(crate) mod tests { server.await.unwrap(); assert_eq!(outcome.unwrap().pages[0].markdown, "native"); assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); assert!(matches!( - machine.resume(None).await, + machine.resume().await, Err(OcrError::InvalidRequest(_)) )); } - async fn drive_native_file_call( - request: crate::ocr::types::LiteLLMOcrRequest, - content: Result, - ) -> (Result, usize) { - let reads = Arc::new(Mutex::new(0)); - let counted = reads.clone(); - let content = Mutex::new(Some(content)); - let host = LocalOcrHost::new(request).with_reader(move || { - *counted.lock().unwrap() += 1; - content.lock().unwrap().take().unwrap() - }); - let outcome = perform_ocr_with(host).await; - let reads = *reads.lock().unwrap(); - (outcome, reads) - } - #[tokio::test] - async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"file"}] - }))]) - .await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::HostReader { - mime_type: Some("application/pdf".into()), - }, - ); - let (response, reads) = drive_native_file_call( - request, - Ok(crate::ocr::types::OcrFileContent { - bytes: b"abc".as_slice().into(), - file_name: Some("scan.png".into()), - }), - ) - .await; - server.await.unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "file"); - assert_eq!(reads, 1); - assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); - } - - #[tokio::test] - async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + async fn empty_byte_documents_fail_before_the_provider_is_called() { let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let failure = OcrError::InvalidRequest("reader exploded".into()); - let (response, reads) = drive_native_file_call( - request - .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Err(failure.clone()), - ) - .await; - assert!( - matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") - ); - assert_eq!(reads, 1); - - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request - .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Ok(crate::ocr::types::OcrFileContent { + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::Bytes { bytes: Default::default(), file_name: None, - }), - ) - .await; + mime_type: None, + }, + ); + let response = perform_ocr_with(LocalOcrHost::new(request)).await; assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); assert!(seen.lock().unwrap().is_empty()); } @@ -3400,24 +3298,21 @@ pub(crate) mod tests { mime_type: None, }, ); - let (response, reads) = - drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; + let (response, ops, _) = drive_until(ocr_client(), &LocalOcrHost::new(request), Ok).await; server.await.unwrap(); std::fs::remove_dir_all(&dir).unwrap(); assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(reads, 0); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::Path { + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::Path { path: path.clone(), mime_type: None, - }), - Err(OcrError::InvalidRequest("unused".into())), - ) - .await; + }, + ); + let response = perform_ocr_with(LocalOcrHost::new(request)).await; assert!(matches!( response.unwrap_err(), OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound @@ -3441,28 +3336,25 @@ pub(crate) mod tests { assert!( matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") ); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); + assert_eq!(ops, ["Project", "BeforeSend"]); + assert!(machine.resume().await.is_err()); } #[tokio::test] - async fn missing_host_result_preserves_pending_operation() { + async fn resuming_before_answering_preserves_pending_operation() { let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); let mut machine = ocr_machine(ocr_client()); + let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else { + panic!("expected the projection op first"); + }; + assert!(machine.resume().await.is_err()); + reply.send(OcrProjection { + request, + caller_token: false, + }); assert!(matches!( - machine.resume(None).await.unwrap(), - MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) - )); - assert!(machine.resume(None).await.is_err()); - assert!(matches!( - machine - .resume(Some(HostResult::Route(OcrOpResult::Request { - request: Box::new(request), - caller_token: false, - }))) - .await - .unwrap(), - MachineStep::Host(HostOp::BeforeSend { .. }) + machine.resume().await, + Ok(MachineStep::Host(HostOp::BeforeSend { .. })) )); } @@ -3623,20 +3515,18 @@ pub(crate) mod tests { }; let host = LocalOcrHost::new(request); let mut machine = ocr_machine(ocr_client()); - let mut result = None; tokio::time::timeout(std::time::Duration::from_secs(2), async { loop { tokio::select! { _ = entered.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { - HostResult::BeforeSend(wire) - } - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + step = machine.resume() => { + match step.unwrap() { + MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), + MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), + MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), + MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), MachineStep::Complete(_) => panic!("pending provider completed"), - }); + } } } } @@ -3661,24 +3551,23 @@ pub(crate) mod tests { } impl Host for CallerTokenHost { - async fn route(&self, op: OcrOp) -> Result { + async fn project(&self) -> Result { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrProjection { + request: self.request.lock().unwrap().take().unwrap(), + caller_token: true, + }) + } + + async fn custom_op(&self, op: OcrOp) -> Result<(), OcrError> { match op { - OcrOp::ProjectRequest => { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrOpResult::Request { - request: Box::new(self.request.lock().unwrap().take().unwrap()), - caller_token: true, - }) - } - OcrOp::AcquireAzureAdToken => { + OcrOp::AcquireAzureAdToken(reply) => { self.trace.lock().unwrap().push("token".into()); - Ok(OcrOpResult::AzureAdToken( - litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( - "caller-token", - )), - )) + reply.send(litellm_auth::ResolvedCredential::Static( + litellm_auth::SecretValue::new("caller-token"), + )); + Ok(()) } - OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), } } @@ -3761,18 +3650,18 @@ pub(crate) mod tests { }); let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); let mut machine = ocr_machine(ocr_client()); - let mut result = None; tokio::time::timeout(std::time::Duration::from_secs(2), async { loop { tokio::select! { _ = received.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + step = machine.resume() => { + match step.unwrap() { + MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), + MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), + MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), + MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), MachineStep::Complete(_) => panic!("the stalled provider completed"), - }); + } } } } diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 59c9cec8da9..20a21e43676 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -25,9 +25,6 @@ pub enum OcrDocumentInput { file_name: Option, mime_type: Option, }, - HostReader { - mime_type: Option, - }, } impl From for OcrDocumentInput { @@ -45,12 +42,6 @@ impl From for OcrDocumentInput { } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct OcrFileContent { - pub bytes: Bytes, - pub file_name: Option, -} - /// Caller-supplied connection overrides for a [`LiteLLMOcrRequest`], in the /// shape hosts receive them: JSON-ish headers, optional timeout, optional /// credentials, and per-field provenance in `input_sources`. diff --git a/litellm-rust/crates/coroutine/AGENTS.md b/litellm-rust/crates/coroutine/AGENTS.md new file mode 100644 index 00000000000..fcb4f4df47f --- /dev/null +++ b/litellm-rust/crates/coroutine/AGENTS.md @@ -0,0 +1,31 @@ +# Requirements + +Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below; the alternatives section says which one each rejected design breaks + +- R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime +- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller; a Rust-side `into_future` would run it in a copied context +- R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime +- R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states +- R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for +- R6 A yield the body makes while being resumed is returned by that same poll, so the host driver's inline first poll needs no extra event-loop turn per op +- R7 No task is spawned: `cancel`, or dropping the coroutine, drops the body, and nothing waits forever on an answer that cannot come +- R8 Several yields can be pending at once, since route code hands clones of its `Co` to token providers and hooks +- R9 Stable Rust + +# Other implementations and why they do not fit + +- Nightly `std::ops::Coroutine`: breaks R9, and its body cannot await futures between yields (R3) +- `genawaiter`: resumes async bodies only with a noop waker, so the body cannot await real I/O (R3) +- `simple_coro`: typestate `Coro` makes answering before resuming a compile-time rule, but its body cannot await arbitrary futures (R3) and its reply type `R` is fixed per coroutine (R5) +- `corosensei` and other stackful coroutines: sync bodies on their own stack, no async I/O inside (R3) +- A hand-written phase enum with an `advance` match (the old `HostPhase`): every await point becomes a state (R4) +- An injected host trait with `async fn`s: core would call the host itself (R1, R2) +- Sans-IO, where core does no I/O and HTTP becomes one more host op: keeps every requirement and makes `resume` a pure step function, but HTTP, streaming, retries and timeouts would move out of core into every bridge; the one real alternative, not taken +- Temporal's Rust workflow SDK (`WorkflowFuture`, `WfContext`) is the closest precedent: an `async fn` polled in place, commands sent over a channel with a oneshot to unblock them. Roles are inverted there (the language SDK owns the program, core answers), and its workflow body may not do real I/O + +# Tradeoffs accepted + +- A tokio `mpsc` channel plus a `oneshot` per yield instead of compiler-generated states +- Protocol mistakes (resuming before answering, resuming after the end) are runtime `ResumeError`s, not compile errors +- Pending yields come out one per `resume`, in the order they were made, and each reply goes back to the yield that made it (R8) +- An answer sent after its yield stopped waiting (for example the body timed out on it) is discarded, since the body already moved on diff --git a/litellm-rust/crates/coroutine/Cargo.toml b/litellm-rust/crates/coroutine/Cargo.toml new file mode 100644 index 00000000000..3ff79ac5f2c --- /dev/null +++ b/litellm-rust/crates/coroutine/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-coroutine" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "Async coroutines on stable Rust whose every yield carries its own typed reply" + +[dependencies] +thiserror.workspace = true +tokio = { workspace = true, features = ["sync"] } + +[dev-dependencies] +rstest.workspace = true +tokio = { workspace = true, features = ["rt", "macros", "time"] } diff --git a/litellm-rust/crates/coroutine/src/co.rs b/litellm-rust/crates/coroutine/src/co.rs new file mode 100644 index 00000000000..d84b041931c --- /dev/null +++ b/litellm-rust/crates/coroutine/src/co.rs @@ -0,0 +1,42 @@ +use std::sync::Weak; + +use tokio::sync::mpsc; + +use crate::{Abandoned, Reply, reply}; + +pub(crate) struct Request { + pub(crate) value: Y, + pub(crate) outstanding: Weak<()>, +} + +/// The body's handle for yielding, `genawaiter`'s `Co`. +pub struct Co { + yields: mpsc::UnboundedSender>, +} + +impl Clone for Co { + fn clone(&self) -> Self { + Self { + yields: self.yields.clone(), + } + } +} + +impl Co { + pub(crate) fn new(yields: mpsc::UnboundedSender>) -> Self { + Self { yields } + } + + /// Yields the value `ask` builds around a fresh [`Reply`] and waits for its answer. + pub async fn yield_(&self, ask: impl FnOnce(Reply) -> Y) -> Result { + let (reply, answer) = reply(); + let outstanding = reply.outstanding(); + self.yields + .send(Request { + value: ask(reply), + outstanding, + }) + .map_err(|_| Abandoned)?; + answer.await + } +} diff --git a/litellm-rust/crates/coroutine/src/coroutine.rs b/litellm-rust/crates/coroutine/src/coroutine.rs new file mode 100644 index 00000000000..fa34f816a37 --- /dev/null +++ b/litellm-rust/crates/coroutine/src/coroutine.rs @@ -0,0 +1,94 @@ +use std::{ + future::{Future, poll_fn}, + pin::Pin, + sync::Weak, + task::{Context, Poll}, +}; + +use tokio::sync::mpsc; + +use crate::{Co, ResumeError, co::Request}; + +/// What one `resume` produced, as in [`std::ops::CoroutineState`]. +#[derive(Debug, PartialEq, Eq)] +pub enum CoroutineState { + Yielded(Y), + Complete(C), +} + +type Body = Pin + Send>>; + +enum Step { + Yielded(Request), + Complete(C), +} + +fn queued( + yields: &mut mpsc::UnboundedReceiver>, + context: &mut Context<'_>, +) -> Option> { + match yields.poll_recv(context) { + Poll::Ready(request) => request, + Poll::Pending => None, + } +} + +pub struct Coroutine { + body: Option>, + yields: mpsc::UnboundedReceiver>, + outstanding: Weak<()>, +} + +impl Coroutine { + /// Builds the body from `producer`. Nothing runs until the first `resume`. + pub fn new(producer: impl FnOnce(Co) -> F) -> Self + where + F: Future + Send + 'static, + { + let (sender, yields) = mpsc::unbounded_channel(); + Self { + body: Some(Box::pin(producer(Co::new(sender)))), + yields, + outstanding: Weak::new(), + } + } + + pub async fn resume(&mut self) -> Result, ResumeError> { + let Some(body) = self.body.as_mut() else { + return Err(ResumeError::Finished); + }; + if self.outstanding.strong_count() > 0 { + return Err(ResumeError::Unanswered); + } + let yields = &mut self.yields; + let step = poll_fn(|context| { + if let Some(request) = queued(yields, context) { + return Poll::Ready(Step::Yielded(request)); + } + if let Poll::Ready(output) = body.as_mut().poll(context) { + return Poll::Ready(Step::Complete(output)); + } + queued(yields, context) + .map_or(Poll::Pending, |request| Poll::Ready(Step::Yielded(request))) + }) + .await; + match step { + Step::Yielded(Request { value, outstanding }) => { + self.outstanding = outstanding; + Ok(CoroutineState::Yielded(value)) + } + Step::Complete(output) => { + self.cancel(); + Ok(CoroutineState::Complete(output)) + } + } + } + + /// Drops the body and fails every yield still waiting, or yet to be made, with + /// [`Abandoned`](crate::Abandoned). + pub fn cancel(&mut self) { + self.body = None; + self.yields.close(); + while self.yields.try_recv().is_ok() {} + } +} diff --git a/litellm-rust/crates/coroutine/src/error.rs b/litellm-rust/crates/coroutine/src/error.rs new file mode 100644 index 00000000000..b8fded23fdb --- /dev/null +++ b/litellm-rust/crates/coroutine/src/error.rs @@ -0,0 +1,14 @@ +/// A `resume` the coroutine refused, leaving it as it was. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum ResumeError { + #[error("coroutine resumed after it finished")] + Finished, + #[error("coroutine resumed before the reply to its last yield was sent or dropped")] + Unanswered, +} + +/// No answer will come to a yield: its [`Reply`](crate::Reply) was dropped unsent, or the +/// coroutine it was sent to is gone. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +#[error("the yield was abandoned before it was answered")] +pub struct Abandoned; diff --git a/litellm-rust/crates/coroutine/src/lib.rs b/litellm-rust/crates/coroutine/src/lib.rs new file mode 100644 index 00000000000..636aaf1b2b6 --- /dev/null +++ b/litellm-rust/crates/coroutine/src/lib.rs @@ -0,0 +1,12 @@ +//! Async coroutines on stable Rust whose every yield carries its own typed [`Reply`]. +//! See `AGENTS.md` for the requirement, the alternatives and the contracts. + +mod co; +mod coroutine; +mod error; +mod reply; + +pub use co::Co; +pub use coroutine::{Coroutine, CoroutineState}; +pub use error::{Abandoned, ResumeError}; +pub use reply::{Answer, Reply, reply}; diff --git a/litellm-rust/crates/coroutine/src/reply.rs b/litellm-rust/crates/coroutine/src/reply.rs new file mode 100644 index 00000000000..b3cb7da2e9d --- /dev/null +++ b/litellm-rust/crates/coroutine/src/reply.rs @@ -0,0 +1,60 @@ +use std::{ + fmt, + future::Future, + pin::Pin, + sync::{Arc, Weak}, + task::{Context, Poll}, +}; + +use tokio::sync::oneshot; + +use crate::Abandoned; + +/// The one way to answer a yield. Sending or dropping it settles the yield. +pub struct Reply { + slot: oneshot::Sender, + outstanding: Arc<()>, +} + +impl Reply { + /// An answer the yield no longer awaits is discarded. + pub fn send(self, answer: A) { + let _ = self.slot.send(answer); + } + + /// Alive until this reply is sent or dropped. + pub(crate) fn outstanding(&self) -> Weak<()> { + Arc::downgrade(&self.outstanding) + } +} + +impl fmt::Debug for Reply { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Reply") + } +} + +/// The waiting end of a [`Reply`]. +pub struct Answer { + slot: oneshot::Receiver, +} + +impl Future for Answer { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { + Pin::new(&mut self.slot) + .poll(context) + .map(|answer| answer.map_err(|_| Abandoned)) + } +} + +/// A reply outside any coroutine, for answering a host operation directly. +pub fn reply() -> (Reply, Answer) { + let (slot, answer) = oneshot::channel(); + let reply = Reply { + slot, + outstanding: Arc::new(()), + }; + (reply, Answer { slot: answer }) +} diff --git a/litellm-rust/crates/coroutine/tests/coroutine.rs b/litellm-rust/crates/coroutine/tests/coroutine.rs new file mode 100644 index 00000000000..91d195503df --- /dev/null +++ b/litellm-rust/crates/coroutine/tests/coroutine.rs @@ -0,0 +1,256 @@ +use std::{ + future::Future, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_coroutine::{Abandoned, Co, Coroutine, CoroutineState, Reply, ResumeError, reply}; +use rstest::rstest; +use tokio::time::timeout; + +#[derive(Debug)] +enum Ask { + Name(Reply<&'static str>), + Count(Reply), +} + +type Test = Coroutine; + +fn yielded(state: Result, ResumeError>) -> Ask { + match state { + Ok(CoroutineState::Yielded(ask)) => ask, + Ok(CoroutineState::Complete(_)) => panic!("expected a yield, the body returned"), + Err(error) => panic!("expected a yield, resume failed: {error}"), + } +} + +fn complete(state: Result, ResumeError>) -> C { + match state { + Ok(CoroutineState::Complete(output)) => output, + Ok(CoroutineState::Yielded(ask)) => panic!("expected completion, got {ask:?}"), + Err(error) => panic!("expected completion, resume failed: {error}"), + } +} + +fn name(ask: Ask) -> Reply<&'static str> { + match ask { + Ask::Name(reply) => reply, + other => panic!("expected a name ask, got {other:?}"), + } +} + +fn count(ask: Ask) -> Reply { + match ask { + Ask::Count(reply) => reply, + other => panic!("expected a count ask, got {other:?}"), + } +} + +/// A body parked at one name ask, with nothing else going on. +fn suspended_once() -> Test> { + Coroutine::new(|co| async move { co.yield_(Ask::Name).await }) +} + +#[tokio::test] +async fn each_typed_answer_resumes_the_yield_that_asked_for_it() { + let mut coroutine: Test = Coroutine::new(|co| async move { + let first = co.yield_(Ask::Name).await.unwrap(); + let second = co.yield_(Ask::Count).await.unwrap(); + format!("{first}+{second}") + }); + + name(yielded(coroutine.resume().await)).send("a"); + count(yielded(coroutine.resume().await)).send(2); + + assert_eq!(complete(coroutine.resume().await), "a+2"); +} + +/// A driver that polls `resume` once, inline, sees every yield the body makes during +/// that poll instead of being sent back to its event loop. +#[test] +fn a_yield_made_while_resuming_is_returned_by_that_same_poll() { + let mut coroutine: Test = Coroutine::new(|co| async move { + let first = co.yield_(Ask::Count).await.unwrap(); + let second = co.yield_(Ask::Count).await.unwrap(); + first + second + }); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + let mut poll_once = + |coroutine: &mut Test| match std::pin::pin!(coroutine.resume()).poll(&mut context) { + std::task::Poll::Ready(state) => state, + std::task::Poll::Pending => panic!("resume needed a second poll"), + }; + + count(yielded(poll_once(&mut coroutine))).send(1); + count(yielded(poll_once(&mut coroutine))).send(2); + + assert_eq!(complete(poll_once(&mut coroutine)), 3); +} + +#[tokio::test] +async fn the_body_awaits_real_futures_between_yields() { + let mut coroutine: Test = Coroutine::new(|co| async move { + tokio::time::sleep(Duration::from_millis(5)).await; + co.yield_(Ask::Count).await.unwrap() + }); + + count(yielded(coroutine.resume().await)).send(7); + + assert_eq!(complete(coroutine.resume().await), 7); +} + +#[tokio::test] +async fn concurrent_yields_come_out_in_order_and_are_answered_separately() { + let mut coroutine: Test<(&str, u32)> = Coroutine::new(|co| async move { + let (first, second) = tokio::join!(co.yield_(Ask::Name), co.yield_(Ask::Count)); + (first.unwrap(), second.unwrap()) + }); + + name(yielded(coroutine.resume().await)).send("one"); + count(yielded(coroutine.resume().await)).send(2); + + assert_eq!(complete(coroutine.resume().await), ("one", 2)); +} + +#[tokio::test] +async fn resuming_before_the_reply_is_settled_is_refused_and_keeps_the_yield_waiting() { + let mut coroutine = suspended_once(); + let reply = name(yielded(coroutine.resume().await)); + + assert_eq!( + coroutine.resume().await.unwrap_err(), + ResumeError::Unanswered + ); + + reply.send("real"); + assert_eq!(complete(coroutine.resume().await), Ok("real")); +} + +#[tokio::test] +async fn a_dropped_reply_abandons_its_yield() { + let mut coroutine = suspended_once(); + drop(yielded(coroutine.resume().await)); + + assert_eq!(complete(coroutine.resume().await), Err(Abandoned)); +} + +#[tokio::test] +async fn an_answer_the_yield_no_longer_awaits_is_discarded() { + let mut coroutine: Test<&str> = Coroutine::new(|co| async move { + tokio::select! { + biased; + _ = co.yield_(Ask::Name) => unreachable!("the answer comes after the body moved on"), + () = std::future::ready(()) => {} + } + co.yield_(Ask::Name).await.unwrap() + }); + let stale = name(yielded(coroutine.resume().await)); + stale.send("stale"); + + name(yielded(coroutine.resume().await)).send("fresh"); + + assert_eq!(complete(coroutine.resume().await), "fresh"); +} + +#[rstest] +#[case::returned(false)] +#[case::cancelled(true)] +#[tokio::test] +async fn a_finished_coroutine_refuses_to_resume(#[case] cancel: bool) { + let mut coroutine = suspended_once(); + let reply = name(yielded(coroutine.resume().await)); + if cancel { + coroutine.cancel(); + } else { + reply.send("done"); + complete(coroutine.resume().await).unwrap(); + } + + assert_eq!(coroutine.resume().await.unwrap_err(), ResumeError::Finished); +} + +#[tokio::test] +async fn a_dropped_resume_leaves_the_coroutine_resumable() { + let mut coroutine: Test = Coroutine::new(|co| async move { + tokio::time::sleep(Duration::from_millis(20)).await; + co.yield_(Ask::Count).await.unwrap() + }); + assert!( + timeout(Duration::from_millis(1), coroutine.resume()) + .await + .is_err() + ); + + count(yielded(coroutine.resume().await)).send(3); + + assert_eq!(complete(coroutine.resume().await), 3); +} + +struct Dropped(Arc>); + +impl Drop for Dropped { + fn drop(&mut self) { + *self.0.lock().unwrap() = true; + } +} + +#[tokio::test] +async fn cancel_drops_the_body() { + let dropped = Arc::new(Mutex::new(false)); + let guard = Dropped(Arc::clone(&dropped)); + let mut coroutine: Test<()> = Coroutine::new(|co| async move { + let _guard = guard; + co.yield_(Ask::Count).await.unwrap(); + }); + let _reply = yielded(coroutine.resume().await); + + coroutine.cancel(); + + assert!(*dropped.lock().unwrap()); +} + +#[rstest] +#[case::cancelled(true)] +#[case::dropped(false)] +#[tokio::test] +async fn a_co_that_escaped_the_body_is_abandoned_once_the_coroutine_ends(#[case] cancel: bool) { + let escaped: Arc>>> = Arc::default(); + let slot = Arc::clone(&escaped); + let mut coroutine: Test<()> = Coroutine::new(move |co| { + *slot.lock().unwrap() = Some(co.clone()); + async move { + co.yield_(Ask::Count).await.unwrap(); + } + }); + let _reply = yielded(coroutine.resume().await); + let co = escaped.lock().unwrap().take().unwrap(); + let waiting = tokio::spawn(async move { co.yield_(Ask::Name).await }); + tokio::task::yield_now().await; + + if cancel { + coroutine.cancel(); + } else { + drop(coroutine); + } + + let outcome = timeout(Duration::from_secs(1), waiting) + .await + .expect("an escaped yield waits forever") + .unwrap(); + assert_eq!(outcome, Err(Abandoned)); +} + +#[rstest] +#[case::sent(true)] +#[case::dropped(false)] +#[tokio::test] +async fn a_detached_reply_settles_its_answer(#[case] send: bool) { + let (reply, answer) = reply::(); + if send { + reply.send(5); + } else { + drop(reply); + } + + assert_eq!(answer.await, if send { Ok(5) } else { Err(Abandoned) }); +} diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 5aca13eeb18..7c1919f9f39 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,8 +1,8 @@ - Target invariants; implementation and runtime validation may lag these rules -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`ProtocolHost` traits - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) + - `ProtocolHost::project` receives the keyword view the adapter's `begin` returned, not the caller's dict; a protocol host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index fb6379dc35a..c1b35c0f69d 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +bytes.workspace = true futures-util.workspace = true litellm-host.workspace = true pyo3.workspace = true diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 3a4cb49be4d..87481aa89b7 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,5 +1,5 @@ use litellm_host::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; -use litellm_host::route::Route; +use litellm_host::protocol::Protocol; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -85,7 +85,7 @@ pub trait PythonLifecycle: Send + Sync { fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -/// Why a route operation the host answered did not produce a result: the route's own code +/// Why a custom operation the host answered did not produce a result: the route's own code /// rejected it, which the route classifies like any other native failure, or Python code /// raised, which reaches the caller as it was raised. #[derive(Debug)] @@ -100,45 +100,54 @@ impl From for InvokeError { } } -/// The Python side of one route: answers the route's own operations, builds the public +/// The Python side of one protocol: answers its custom operations, builds the public /// response and classifies native failures into public exceptions. -pub trait RouteHost: Send + Sync { - type Route: Route; +pub trait ProtocolHost: Send + Sync { + type Protocol: Protocol; /// The public exception a native failure maps to, kept as a value until the driver /// raises it. type Failure: Into; - /// `arguments` is the keyword view the lifecycle's `begin` produced, not the - /// caller's own dict. A route host that projects from it inherits whatever that - /// adapter rewrote. - fn invoke( + /// Projects the call's request. `arguments` is the keyword view the lifecycle's + /// `begin` produced, not the caller's own dict, so the projection inherits whatever + /// that adapter rewrote. + fn project( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - op: ::Op, - ) -> Result<::OpResult, InvokeError<::Error>>; + ) -> Result< + ::Projection, + InvokeError<::Error>, + >; + + /// Answers `op` through its reply. + fn invoke( + &mut self, + py: Python<'_>, + op: ::Op, + ) -> Result<(), InvokeError<::Error>>; fn complete( &mut self, py: Python<'_>, - response: ::Response, + response: ::Response, ) -> PyResult>; /// One streamed chunk as the caller receives it. fn chunk( &mut self, py: Python<'_>, - chunk: ::Chunk, + chunk: ::Chunk, ) -> PyResult>; fn classify( &self, py: Python<'_>, - error: ::Error, + error: ::Error, ) -> PyResult; - fn host_error(error: &PyErr) -> ::Error; + fn host_error(error: &PyErr) -> ::Error; fn close(&mut self, py: Python<'_>); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 77a294d274b..aaa0752522b 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,10 +2,11 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; +use litellm_host::event::WireRequest; use litellm_host::event::{FailureOrigin, Timing, epoch_seconds}; -use litellm_host::host::{Demand, HostOp, HostResult, HostStep}; +use litellm_host::host::{Demand, HostOp, HostStep, Reply}; use litellm_host::machine::{HostFailure, Machine, MachineStep}; -use litellm_host::route::Route; +use litellm_host::protocol::Protocol; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -13,21 +14,21 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; -type RouteOf = ::Route; -type ErrorOf = as Route>::Error; -type ResponseOf = as Route>::Response; -type NativeStep = MachineStep, ResponseOf>; +type ProtocolOf = ::Protocol; +type ErrorOf = as Protocol>::Error; +type ResponseOf = as Protocol>::Response; +type NativeStep = MachineStep, ResponseOf>; type NativeResult = Result, ErrorOf>; -type NativeResume = Option>, HostFailure>>>; +type Interruption = Option>>; type MachineResult = Result< - MachineStep<::Route, ::Complete>, - <::Route as Route>::Error, + MachineStep<::Protocol, ::Complete>, + <::Protocol as Protocol>::Error, >; struct MachineState { @@ -44,12 +45,11 @@ enum Stage { Failed(Py), } -#[derive(Clone, Copy)] enum Expect { Started, Arguments, - Wire, - Emitted, + Wire(Reply), + Emitted(Reply<()>), Response, Terminal, } @@ -58,20 +58,30 @@ enum Pending { Native, Adapter(Expect), /// The stream handed to the caller waits for its next read or its close. - Consumer, + Consumer(Reply), } -enum Next { +/// A route answer as the driver resumes on it: a Python exception interrupts the call as +/// raised, a native rejection resumes the machine with it. +fn answered(answer: Result<(), InvokeError>) -> PyResult> { + match answer { + Ok(()) => Ok(Ok(())), + Err(InvokeError::Native(error)) => Ok(Err(error)), + Err(InvokeError::Python(error)) => Err(error), + } +} + +enum Next { Return(ExecutionStep), Continue(HostStep, Py>), } struct PythonDriver where - H: RouteHost, - M: Machine> + 'static, + H: ProtocolHost, + M: Machine> + 'static, { - route: H, + host: H, adapter: Box, machine: Option>>>, arguments: Option>, @@ -89,17 +99,17 @@ where pub fn run_call( py: Python<'_>, machine: M, - route: H, + host: H, adapter: Box, arguments: Py, asynchronous: bool, ) -> PyResult> where - H: RouteHost + 'static, - M: Machine> + 'static, + H: ProtocolHost + 'static, + M: Machine> + 'static, { let mut driver = PythonDriver { - route, + host, adapter, machine: Some(Arc::new(Mutex::new(MachineState { machine, @@ -141,8 +151,8 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { impl PythonDriver where - H: RouteHost, - M: Machine> + 'static, + H: ProtocolHost, + M: Machine> + 'static, { fn timing(&self) -> Timing { Timing { @@ -172,13 +182,13 @@ where self.run_steps(py, HostStep::Ready(result)) } (Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error), - (Some(Pending::Consumer), Some(read)) => { - let demand = if read.is_ok() { + (Some(Pending::Consumer(reply)), Some(read)) => { + reply.send(if read.is_ok() { Demand::More } else { Demand::Detached - }; - self.resume_machine(py, Some(Ok(HostResult::Demand(demand)))) + }); + self.resume_machine(py, None) } (Some(Pending::Adapter(expect)), Some(result)) => { match self.adapter.resume(py, result) { @@ -196,22 +206,24 @@ where step: LifecycleStep, expect: Expect, ) -> PyResult { + if let LifecycleStep::Await(awaitable) = step { + self.pending = Some(Pending::Adapter(expect)); + return Ok(ExecutionStep::Await(awaitable)); + } match (expect, step) { - (_, LifecycleStep::Await(awaitable)) => { - self.pending = Some(Pending::Adapter(expect)); - Ok(ExecutionStep::Await(awaitable)) - } (Expect::Started, LifecycleStep::Done) => self.begin(py), (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) } - (Expect::Wire, LifecycleStep::Wire(wire)) => { - self.resume_machine(py, Some(Ok(HostResult::BeforeSend(wire)))) + (Expect::Wire(reply), LifecycleStep::Wire(wire)) => { + reply.send(*wire); + self.resume_machine(py, None) } - (Expect::Emitted, LifecycleStep::Done) => { - self.resume_machine(py, Some(Ok(HostResult::Emitted))) + (Expect::Emitted(reply), LifecycleStep::Done) => { + reply.send(()); + self.resume_machine(py, None) } (Expect::Response, LifecycleStep::Response(response)) => self.succeeded(py, response), (Expect::Terminal, LifecycleStep::Done) => match &self.stage { @@ -242,9 +254,9 @@ where fn resume_machine( &mut self, py: Python<'_>, - result: NativeResume, + interruption: Interruption, ) -> PyResult { - let step = self.resume_core(py, result)?; + let step = self.resume_core(py, interruption)?; self.run_steps(py, step) } @@ -277,53 +289,62 @@ where } Err(error) => return self.machine_failed(py, error).map(Next::Return), }; - let answer = match op { - HostOp::Route(op) => { + let answered = match op { + HostOp::Project(reply) => { let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; - match self.route.invoke(py, arguments.bind(py), op) { - Ok(result) => Ok(HostResult::Route(result)), - Err(InvokeError::Native(error)) => { - return self - .resume_core(py, Some(Err(HostFailure::Error(error)))) - .map(Next::Continue); - } - Err(InvokeError::Python(error)) => Err(error), - } + let projected = self.host.project(py, arguments.bind(py)); + answered(projected.map(|projection| reply.send(projection))) } - HostOp::BeforeSend { wire, context } => { - match self.adapter.before_send(py, wire, &context) { - Ok(LifecycleStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), + HostOp::Custom(op) => answered(self.host.invoke(py, op)), + HostOp::BeforeSend { + wire, + context, + reply, + } => match self.adapter.before_send(py, wire, &context) { + Ok(LifecycleStep::Wire(wire)) => { + reply.send(*wire); + Ok(Ok(())) + } + Ok(LifecycleStep::Await(awaitable)) => { + self.pending = Some(Pending::Adapter(Expect::Wire(reply))); + return Ok(Next::Return(ExecutionStep::Await(awaitable))); + } + Ok(_) => return Err(missing_state()), + Err(error) => Err(error), + }, + HostOp::Open(_, reply) => return self.opened(py, reply).map(Next::Return), + HostOp::Deliver(chunk, reply) => { + return self.delivered(py, chunk, reply).map(Next::Return); + } + HostOp::Emit(event, reply) => { + match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { + Ok(LifecycleStep::Done) => { + reply.send(()); + Ok(Ok(())) + } Ok(LifecycleStep::Await(awaitable)) => { - self.pending = Some(Pending::Adapter(Expect::Wire)); + self.pending = Some(Pending::Adapter(Expect::Emitted(reply))); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } Ok(_) => return Err(missing_state()), Err(error) => Err(error), } } - HostOp::Open(_) => return self.opened(py).map(Next::Return), - HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), - HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { - Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), - Ok(LifecycleStep::Await(awaitable)) => { - self.pending = Some(Pending::Adapter(Expect::Emitted)); - return Ok(Next::Return(ExecutionStep::Await(awaitable))); - } - Ok(_) => return Err(missing_state()), - Err(error) => Err(error), - }, }; - match answer { - Ok(answer) => self.resume_core(py, Some(Ok(answer))).map(Next::Continue), + match answered { + Ok(Ok(())) => self.resume_core(py, None).map(Next::Continue), + Ok(Err(native)) => self + .resume_core(py, Some(HostFailure::Error(native))) + .map(Next::Continue), Err(error) => self.interrupt(py, error).map(Next::Return), } } - fn opened(&mut self, py: Python<'_>) -> PyResult { + fn opened(&mut self, py: Python<'_>, reply: Reply) -> PyResult { self.stage = Stage::Streaming; match self.adapter.opened(py) { Ok(()) => { - self.pending = Some(Pending::Consumer); + self.pending = Some(Pending::Consumer(reply)); Ok(ExecutionStep::Open) } Err(error) => self.interrupt(py, error), @@ -333,15 +354,16 @@ where fn delivered( &mut self, py: Python<'_>, - chunk: as Route>::Chunk, + chunk: as Protocol>::Chunk, + reply: Reply, ) -> PyResult { - let chunk = match self.route.chunk(py, chunk) { + let chunk = match self.host.chunk(py, chunk) { Ok(chunk) => chunk, Err(error) => return self.interrupt(py, error), }; match self.adapter.delivered(py, &chunk) { Ok(()) => { - self.pending = Some(Pending::Consumer); + self.pending = Some(Pending::Consumer(reply)); Ok(ExecutionStep::Yield(chunk)) } Err(error) => self.interrupt(py, error), @@ -357,25 +379,24 @@ where } else { HostFailure::Error(native) }; - self.resume_machine(py, Some(Err(failure))) + self.resume_machine(py, Some(failure)) } fn resume_core( &mut self, py: Python<'_>, - result: NativeResume, + interruption: Interruption, ) -> PyResult, Py>> { let state = Arc::clone(self.machine.as_ref().ok_or_else(missing_state)?); let future = async move { let mut state = state.lock().await; - let result = match result { - Some(Err(failure)) => state + let result = match interruption { + Some(failure) => state .machine .interrupt(failure) .await .map(MachineStep::Complete), - Some(Ok(result)) => state.machine.resume(Some(result)).await, - None => state.machine.resume(None).await, + None => state.machine.resume().await, }; state.result = Some(result); Ok(()) @@ -414,7 +435,7 @@ where fn completed(&mut self, py: Python<'_>, response: ResponseOf) -> PyResult { self.ended_at = Some(epoch_seconds()); - let public = match self.route.complete(py, response) { + let public = match self.host.complete(py, response) { Ok(public) => public, Err(error) => return self.failure(py, error, FailureOrigin::Call), }; @@ -441,7 +462,7 @@ where /// fails, that failure is raised with the native error's text as its `__context__`. fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { let native = error.to_string(); - let classifier_error = match self.route.classify(py, error) { + let classifier_error = match self.host.classify(py, error) { Ok(failure) => return failure.into(), Err(classifier_error) => classifier_error, }; @@ -486,7 +507,7 @@ where if self.machine.take().is_some() { Python::attach(|py| { self.adapter.close(py); - self.route.close(py); + self.host.close(py); }); } } @@ -494,15 +515,15 @@ where impl ExecutionBody for PythonDriver where - H: RouteHost, - M: Machine> + 'static, + H: ProtocolHost, + M: Machine> + 'static, { fn resume(&mut self, result: Option>>) -> PyResult { Python::attach(|py| self.drive(py, result)) } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.route.traverse(visit)?; + self.host.traverse(visit)?; self.adapter.traverse(visit)?; visit.call(&self.arguments)?; visit.call(&self.interrupted)?; @@ -516,8 +537,8 @@ where impl Drop for PythonDriver where - H: RouteHost, - M: Machine> + 'static, + H: ProtocolHost, + M: Machine> + 'static, { fn drop(&mut self) { self.clear(); @@ -528,8 +549,8 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_host::event::{MachineEvent, RequestContext, WireRequest}; - use litellm_host::machine::{Interrupted, Step}; + use litellm_host::event::{MachineEvent, RawResponse, RequestContext}; + use litellm_host::machine::{CallMachine, MachineFault}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -573,22 +594,21 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } - struct Synthetic; - - impl Route for Synthetic { - type Response = String; - type Error = Error; - type Op = &'static str; - type OpResult = String; - type Chunk = std::convert::Infallible; - type StreamHead = std::convert::Infallible; + impl From for Error { + fn from(fault: MachineFault) -> Self { + Self(format!("{fault:?}")) + } } - /// Yields the scripted ops in order, then completes or fails as scripted. - struct ScriptedMachine { - ops: Vec>, - outcome: Option>, - answers: Vec, + struct Synthetic; + + impl Protocol for Synthetic { + type Response = String; + type Error = Error; + type Projection = String; + type Op = (&'static str, Reply); + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } fn wire() -> WireRequest { @@ -609,37 +629,6 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } - impl Machine for ScriptedMachine { - type Route = Synthetic; - type Complete = String; - - fn resume(&mut self, result: Option>) -> Step<'_, Self> { - Box::pin(async move { - if let Some(result) = result { - self.answers.push(match result { - HostResult::Route(value) => value, - HostResult::BeforeSend(wire) => wire.url, - HostResult::Emitted => "emitted".into(), - HostResult::Demand(demand) => format!("{demand:?}"), - }); - } - if !self.ops.is_empty() { - return Ok(MachineStep::Host(self.ops.remove(0))); - } - self.outcome - .take() - .ok_or_else(|| Error("resumed after completion".into()))? - .map(MachineStep::Complete) - }) - } - - fn interrupt(&mut self, failure: HostFailure) -> Interrupted<'_, Self> { - self.ops.clear(); - self.outcome = None; - Box::pin(async move { Err(failure.into_error()) }) - } - } - #[derive(Default)] struct Log(Arc>>); @@ -677,22 +666,37 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } - impl RouteHost for SyntheticHost { - type Route = Synthetic; + impl SyntheticHost { + fn answer(&self, value: impl FnOnce() -> String) -> Result> { + match self.op { + OpScript::Answer => Ok(value()), + OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), + OpScript::RejectNatively => Err(InvokeError::Native(Error("op rejected".into()))), + } + } + } + + impl ProtocolHost for SyntheticHost { + type Protocol = Synthetic; type Failure = Classified; + fn project( + &mut self, + _: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> Result> { + self.log.push("project"); + self.answer(|| format!("project:{}", arguments.len())) + } + fn invoke( &mut self, _: Python<'_>, - arguments: &Bound<'_, PyDict>, - op: &'static str, - ) -> Result> { - self.log.push(format!("route:{op}")); - match self.op { - OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), - OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), - OpScript::RejectNatively => Err(InvokeError::Native(Error("op rejected".into()))), - } + (op, reply): (&'static str, Reply), + ) -> Result<(), InvokeError> { + self.log.push(format!("op:{op}")); + self.answer(|| op.to_string()) + .map(|answer| reply.send(answer)) } fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { @@ -719,7 +723,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } fn close(&mut self, _: Python<'_>) { - self.log.push("route.close"); + self.log.push("host.close"); } fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -828,7 +832,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_scripted( py: Python<'_>, - machine: ScriptedMachine, + machine: CallMachine, op: OpScript, script: AdapterScript, asynchronous: bool, @@ -848,12 +852,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_hosted( py: Python<'_>, - machine: ScriptedMachine, - route: SyntheticHost, + machine: CallMachine, + host: SyntheticHost, script: AdapterScript, asynchronous: bool, ) -> (PyResult>, Vec) { - let log = Log(route.log.0.clone()); + let log = Log(host.log.0.clone()); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script, @@ -863,7 +867,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let result = run_call( py, machine, - route, + host, Box::new(adapter), arguments.unbind(), asynchronous, @@ -884,21 +888,21 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri (result, log.entries()) } - fn success_machine() -> ScriptedMachine { - ScriptedMachine { - ops: vec![ - HostOp::Route("project"), - HostOp::BeforeSend { - wire: Box::new(wire()), - context: Box::new(context()), - }, - HostOp::Emit(MachineEvent::ResponseReceived { - raw: litellm_host::event::RawResponse { body: "raw".into() }, - }), - ], - outcome: Some(Ok("done".into())), - answers: Vec::new(), - } + /// Answers to projection, to the route op and to `before_send` all reach the + /// response, so a driver that misroutes a reply changes what the call returns. + fn success_machine() -> CallMachine { + CallMachine::new(|host| { + Box::pin(async move { + let projected = host.project().await?; + let signed = host.custom_op(|reply| ("sign", reply)).await?; + let wire = host.before_send(wire(), context()).await?; + host.emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: "raw".into() }, + }) + .await?; + Ok(format!("{projected}|{signed}|{}", wire.url)) + }) + }) } #[test] @@ -917,32 +921,37 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri AdapterScript::Plain, asynchronous, ); - assert_eq!(result.unwrap().extract::(py).unwrap(), "done"); + assert_eq!( + result.unwrap().extract::(py).unwrap(), + "project:1|sign|rewritten" + ); assert_eq!( log, [ "started", "begin", - "route:project", + "project", + "op:sign", "before_send", "response:raw", "complete", "after_success", - "succeeded:done", + "succeeded:project:1|sign|rewritten", "adapter.close", - "route.close", + "host.close", ] ); } }); } - fn failing_machine() -> ScriptedMachine { - ScriptedMachine { - ops: vec![HostOp::Route("project")], - outcome: Some(Err(Error("provider exploded".into()))), - answers: Vec::new(), - } + fn failing_machine() -> CallMachine { + CallMachine::new(|host| { + Box::pin(async move { + host.project().await?; + Err(Error("provider exploded".into())) + }) + }) } #[test] @@ -969,11 +978,11 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri [ "started", "begin", - "route:project", + "project", "classify:provider exploded", "failed:Call:classified: provider exploded", "adapter.close", - "route.close", + "host.close", ] ); } @@ -1003,11 +1012,11 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri [ "started", "begin", - "route:project", + "project", "classify:op rejected", "failed:Call:classified: op rejected", "adapter.close", - "route.close", + "host.close", ] ); }); @@ -1035,10 +1044,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri [ "started", "begin", - "route:project", + "project", "failed:Call:op failed", "adapter.close", - "route.close", + "host.close", ] ); }); @@ -1073,11 +1082,11 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri [ "started", "begin", - "route:project", + "project", "classify:provider exploded", "failed:Call:classifier failed", "adapter.close", - "route.close", + "host.close", ] ); }); @@ -1106,7 +1115,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri "begin", "failed:Host:begin failed", "adapter.close", - "route.close" + "host.close" ] ); }); @@ -1130,7 +1139,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ); assert_eq!(result.unwrap().extract::(py).unwrap(), "replaced"); assert!(log.contains(&"succeeded:replaced".to_string())); - assert!(!log.contains(&"succeeded:done".to_string())); + assert!(!log.contains(&"succeeded:project:1|rewritten".to_string())); } }); } @@ -1159,7 +1168,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri "after_success", "failed:Host:after_success failed", "adapter.close", - "route.close" + "host.close" ] ); assert!(!log.iter().any(|entry| entry.starts_with("succeeded"))); @@ -1175,18 +1184,24 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri crate::initialize_python(); Python::attach(|py| { struct Cancelling(Log); - impl RouteHost for Cancelling { - type Route = Synthetic; + impl ProtocolHost for Cancelling { + type Protocol = Synthetic; type Failure = Classified; - fn invoke( + fn project( &mut self, _: Python<'_>, _: &Bound<'_, PyDict>, - _: &'static str, ) -> Result> { - self.0.push("route"); + self.0.push("project"); Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into()) } + fn invoke( + &mut self, + _: Python<'_>, + _: (&'static str, Reply), + ) -> Result<(), InvokeError> { + Err(missing_state().into()) + } fn chunk( &mut self, _: Python<'_>, @@ -1210,7 +1225,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } let log = Log::default(); - let route = Cancelling(Log(log.0.clone())); + let host = Cancelling(Log(log.0.clone())); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script: AdapterScript::Plain, @@ -1218,7 +1233,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let error = run_call( py, success_machine(), - route, + host, Box::new(adapter), PyDict::new(py).unbind(), false, @@ -1227,7 +1242,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert!(!error.is_instance_of::(py)); assert_eq!( log.entries(), - ["started", "begin", "route", "adapter.close"] + ["started", "begin", "project", "adapter.close"] ); }); } diff --git a/litellm-rust/crates/host-python/src/file_reader.rs b/litellm-rust/crates/host-python/src/file_reader.rs new file mode 100644 index 00000000000..bbc7a233b28 --- /dev/null +++ b/litellm-rust/crates/host-python/src/file_reader.rs @@ -0,0 +1,241 @@ +//! A caller's file-like object: anything with a callable `read`, kept as a handle and read +//! once, on the host's thread, into bytes Rust owns. + +use bytes::Bytes; +use pyo3::{ + exceptions::PyTypeError, + gc::{PyTraverseError, PyVisit}, + prelude::*, + pybacked::PyBackedBytes, + types::{PyBytes, PyString}, +}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FileContent { + pub bytes: Bytes, + pub file_name: Option, +} + +#[derive(Debug)] +pub struct PythonFileReader { + reader: Py, + name: Option, +} + +impl PythonFileReader { + /// `None` when `file` has no callable `read`. The object's `name` is read now, its + /// contents only on [`read`](Self::read). + pub fn from_file_like(file: &Bound<'_, PyAny>) -> PyResult> { + let reader = file + .getattr_opt("read")? + .filter(|value| value.is_callable()); + let Some(reader) = reader else { + return Ok(None); + }; + let name = file + .getattr_opt("name")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; + Ok(Some(Self { + reader: reader.unbind(), + name, + })) + } + + pub fn read(&self, py: Python<'_>) -> PyResult { + let value = self.reader.bind(py).call0()?; + let bytes = if value.is_instance_of::() { + Bytes::from(value.extract::()?) + } else if value.is_instance_of::() { + py_bytes(&value)? + } else { + return Err(PyTypeError::new_err(format!( + "file read must return bytes or str, got {}", + value.get_type(), + ))); + }; + Ok(FileContent { + bytes, + file_name: self.name.clone(), + }) + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reader) + } +} + +/// An exact `bytes` object is shared without copying and keeps the Python object alive; +/// a `bytes` subclass is copied. +pub fn py_bytes(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_exact_instance_of::() { + return Ok(Bytes::from_owner(value.extract::()?)); + } + Ok(Bytes::copy_from_slice( + value.extract::()?.as_ref(), + )) +} + +#[cfg(test)] +mod tests { + use pyo3::{exceptions::PyTypeError, types::PyDict}; + + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn reader<'py>(locals: &Bound<'py, PyDict>, name: &str) -> PythonFileReader { + PythonFileReader::from_file_like(&locals.get_item(name).unwrap().unwrap()) + .unwrap() + .unwrap() + } + + #[test] + fn objects_without_a_callable_read_are_not_readers() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Attribute: + read = 'not callable' +plain = object() +attribute = Attribute() +", + ); + for name in ["plain", "attribute"] { + let file = locals.get_item(name).unwrap().unwrap(); + assert!(PythonFileReader::from_file_like(&file).unwrap().is_none()); + } + }); + } + + #[test] + fn the_name_is_taken_up_front_and_the_contents_only_on_read() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Reader: + name = 'scan.png' + def __init__(self): + self.reads = 0 + def read(self): + self.reads += 1 + return b'abc' +file = Reader() +", + ); + let reads = || { + locals + .get_item("file") + .unwrap() + .unwrap() + .getattr("reads") + .unwrap() + .extract::() + .unwrap() + }; + let file = reader(&locals, "file"); + assert_eq!(reads(), 0); + let content = file.read(py).unwrap(); + assert_eq!(reads(), 1); + assert_eq!( + content, + FileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + } + ); + }); + } + + #[test] + fn read_results_are_normalized_and_exceptions_keep_their_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = KeyError('reader failed') +class Raising: + def read(self): + raise failure +class Text: + def read(self): + return 'héllo' +class Wrong: + def read(self): + return 7 +raising = Raising() +text = Text() +wrong = Wrong() +", + ); + let error = reader(&locals, "raising").read(py).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert_eq!( + reader(&locals, "text").read(py).unwrap().bytes.as_ref(), + "héllo".as_bytes() + ); + let error = reader(&locals, "wrong").read(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("bytes or str")); + }); + } + + #[rstest::rstest] + #[case::read("read")] + #[case::name("name")] + fn attribute_failures_keep_their_identity(#[case] attribute: &str) { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = LookupError('file property failed') +class File: + def __getattribute__(self, name): + if name == attribute: + raise failure + return super().__getattribute__(name) + name = 'scan.pdf' + def read(self): + return b'abc' +file = File() +", + ); + locals.set_item("attribute", attribute).unwrap(); + let error = + PythonFileReader::from_file_like(&locals.get_item("file").unwrap().unwrap()) + .unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() { + Python::initialize(); + let (bytes, pointer) = Python::attach(|py| { + let value = PyBytes::new(py, b"document bytes"); + let pointer = value.as_bytes().as_ptr() as usize; + (py_bytes(value.as_any()).unwrap(), pointer) + }); + assert_eq!(bytes.as_ptr() as usize, pointer); + assert_eq!(bytes.as_ref(), b"document bytes"); + } +} diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 55b27e34b46..7e17c4da51e 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,6 +1,6 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and //! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine) -//! against a Python route host and a Python lifecycle. Everything here is Python-specific by +//! against a Python protocol host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod adapter; @@ -8,13 +8,14 @@ mod argument; mod callable; mod driver; mod execution; +mod file_reader; mod fork_gate; mod gil; mod handle; mod marshal; pub use adapter::{ - InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; @@ -24,6 +25,7 @@ pub use execution::{ reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value, runtime_started, }; +pub use file_reader::{FileContent, PythonFileReader, py_bytes}; pub use fork_gate::RuntimeAlreadyStarted; pub use gil::{PythonContext, attach_blocking, release_count, release_gil}; pub use handle::{Execution, ExecutionBody, ExecutionStep}; diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index 0c7c46192b5..bbbed68f345 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-auth.workspace = true +litellm-coroutine.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/host/src/host.rs b/litellm-rust/crates/host/src/host.rs index aba35185a18..9714b9470a3 100644 --- a/litellm-rust/crates/host/src/host.rs +++ b/litellm-rust/crates/host/src/host.rs @@ -1,28 +1,27 @@ use std::future::Future; -use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; -use crate::route::Route; +pub use litellm_coroutine::{Abandoned, Answer, Reply, reply}; -/// One suspension point of a native call, performed by the host. -pub enum HostOp { - Route(R::Op), +use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; +use crate::protocol::Protocol; + +/// One suspension point of a native call, performed by the host and answered through the +/// [`Reply`] it carries. +pub enum HostOp { + /// The first op of every call: the caller's request as the host projects it. + Project(Reply), + Custom(R::Op), BeforeSend { wire: Box, context: Box, + reply: Reply, }, - Emit(MachineEvent), + Emit(MachineEvent, Reply<()>), /// The response streams: the host hands the caller a stream and answers once the /// caller asks for the first chunk or goes away. - Open(R::StreamHead), + Open(R::StreamHead, Reply), /// The next chunk of an open stream, answered once the caller asks for the one after. - Deliver(R::Chunk), -} - -pub enum HostResult { - Route(R::OpResult), - BeforeSend(Box), - Emitted, - Demand(Demand), + Deliver(R::Chunk, Reply), } /// Whether the caller of a streamed call still reads it. @@ -39,10 +38,13 @@ pub enum HostStep { Suspend(S), } -/// An in-process host: answers route operations and observes the call without leaving +/// An in-process host: answers custom operations and observes the call without leaving /// the Rust runtime. Language hosts implement their own driver instead. -pub trait Host: Send + Sync { - fn route(&self, op: R::Op) -> impl Future> + Send; +pub trait Host: Send + Sync { + fn project(&self) -> impl Future> + Send; + + /// Answers `op` through its reply, or fails the call. + fn custom_op(&self, op: R::Op) -> impl Future> + Send; fn before_send( &self, diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index 65479c2380f..c6b9e59b65a 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -1,12 +1,13 @@ //! The contract between a native call and the host runtime that drives it. //! //! A host is whatever sits on the far side of the language boundary: CPython today, -//! another runtime later. Core runs each route on a [`machine::RouteMachine`] and never learns +//! another runtime later. Core runs each route on a [`machine::CallMachine`] and never learns //! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers -//! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent. +//! each through the typed [`host::Reply`] it carries, observes [`event::CallEvent`]s and +//! may rewrite the wire request before it is sent. pub mod event; pub mod host; pub mod machine; -pub mod route; +pub mod protocol; pub mod run; diff --git a/litellm-rust/crates/host/src/machine/auth.rs b/litellm-rust/crates/host/src/machine/auth.rs index ba7e242e766..76e3504ca28 100644 --- a/litellm-rust/crates/host/src/machine/auth.rs +++ b/litellm-rust/crates/host/src/machine/auth.rs @@ -1,22 +1,21 @@ use std::sync::Arc; use super::{HostChannel, MachineFault}; -use crate::route::Route; +use crate::{host::Reply, protocol::Protocol}; use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -/// A route whose host can mint credentials on the call's behalf. -pub trait TokenRoute: Route { - fn acquire_token_op() -> Self::Op; - fn token_credential(result: Self::OpResult) -> Option; +/// A protocol whose host can mint credentials on the call's behalf. +pub trait TokenProtocol: Protocol { + fn acquire_token_op(reply: Reply) -> Self::Op; } /// A [`TokenProvider`] that asks the host for each credential through the call's own /// operation channel, so the host answers it on the caller's thread and context. -pub struct HostTokenProvider { +pub struct HostTokenProvider { channel: HostChannel, } -impl std::fmt::Debug for HostTokenProvider { +impl std::fmt::Debug for HostTokenProvider { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter.write_str("HostTokenProvider") } @@ -24,7 +23,7 @@ impl std::fmt::Debug for HostTokenProvider { impl HostTokenProvider where - R: TokenRoute, + R: TokenProtocol, R::Error: From + std::fmt::Display, { pub fn handle(channel: HostChannel) -> TokenProviderHandle { @@ -34,19 +33,15 @@ where impl TokenProvider for HostTokenProvider where - R: TokenRoute, + R: TokenProtocol, R::Error: From + std::fmt::Display, { fn acquire(&self) -> TokenFuture<'_> { Box::pin(async move { - let result = self - .channel - .route(R::acquire_token_op()) + self.channel + .custom_op(R::acquire_token_op) .await - .map_err(|error| Error::AzureTokenAcquisition(error.to_string()))?; - R::token_credential(result).ok_or_else(|| { - Error::AzureTokenAcquisition("invalid token provider host result".into()) - }) + .map_err(|error| Error::AzureTokenAcquisition(error.to_string())) }) } } diff --git a/litellm-rust/crates/host/src/machine/call_machine.rs b/litellm-rust/crates/host/src/machine/call_machine.rs new file mode 100644 index 00000000000..af0bc50fbe6 --- /dev/null +++ b/litellm-rust/crates/host/src/machine/call_machine.rs @@ -0,0 +1,137 @@ +//! The one machine every route runs on: the route's provider future as a +//! [`Coroutine`] that yields [`HostOp`]s, each answered through its own typed reply. No +//! task is spawned; dropping the machine drops the in-flight call. + +use std::{future::Future, pin::Pin}; + +use litellm_coroutine::{Co, Coroutine, CoroutineState, ResumeError}; + +use super::{HostFailure, Interrupted, Machine, MachineStep, Step}; +use crate::{ + event::{MachineEvent, RequestContext, WireRequest}, + host::{Demand, HostOp, Reply}, + protocol::Protocol, +}; + +/// The machine's own failures, distinct from anything the provider call reports. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MachineFault { + /// The host dropped an op's reply unanswered, or went away while the call waited. + Abandoned, + /// The host resumed the call out of turn. + Protocol(ResumeError), +} + +pub type ExecuteFuture = + Pin::Response, ::Error>> + Send>>; + +/// The provider side of the machine: how the in-flight call reaches its host. +pub struct HostChannel { + co: Co>, +} + +impl Clone for HostChannel { + fn clone(&self) -> Self { + Self { + co: self.co.clone(), + } + } +} + +impl HostChannel +where + R::Error: From, +{ + async fn yield_( + &self, + ask: impl FnOnce(Reply) -> HostOp + Send, + ) -> Result { + self.co + .yield_(ask) + .await + .map_err(|_| MachineFault::Abandoned.into()) + } + + pub async fn project(&self) -> Result { + self.yield_(HostOp::Project).await + } + + /// Asks the host to perform the custom operation `ask` builds around its reply, as in + /// `host.custom_op(OcrOp::AcquireAzureAdToken)`. + pub async fn custom_op( + &self, + ask: impl FnOnce(Reply) -> R::Op + Send, + ) -> Result { + self.yield_(|reply| HostOp::Custom(ask(reply))).await + } + + pub async fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + self.yield_(|reply| HostOp::BeforeSend { + wire: Box::new(wire), + context: Box::new(context), + reply, + }) + .await + } + + pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { + self.yield_(|reply| HostOp::Emit(event, reply)).await + } + + pub async fn open(&self, head: R::StreamHead) -> Result { + self.yield_(|reply| HostOp::Open(head, reply)).await + } + + pub async fn deliver(&self, chunk: R::Chunk) -> Result { + self.yield_(|reply| HostOp::Deliver(chunk, reply)).await + } +} + +type CallCoroutine = + Coroutine, Result<::Response, ::Error>>; + +pub struct CallMachine { + coroutine: CallCoroutine, +} + +impl CallMachine +where + R::Error: From, +{ + pub fn new(execute: impl FnOnce(HostChannel) -> ExecuteFuture + Send + 'static) -> Self { + Self { + coroutine: Coroutine::new(|co| execute(HostChannel { co })), + } + } +} + +impl Machine for CallMachine +where + R::Error: From, +{ + type Protocol = R; + type Complete = R::Response; + + fn resume(&mut self) -> Step<'_, Self> { + Box::pin(async move { + match self + .coroutine + .resume() + .await + .map_err(MachineFault::Protocol)? + { + CoroutineState::Yielded(op) => Ok(MachineStep::Host(op)), + CoroutineState::Complete(outcome) => outcome.map(MachineStep::Complete), + } + }) + } + + fn interrupt(&mut self, failure: HostFailure) -> Interrupted<'_, Self> { + self.coroutine.cancel(); + Box::pin(async move { Err(failure.into_error()) }) + } +} diff --git a/litellm-rust/crates/host/src/machine/mod.rs b/litellm-rust/crates/host/src/machine/mod.rs index 2c26db61582..0c7501633fa 100644 --- a/litellm-rust/crates/host/src/machine/mod.rs +++ b/litellm-rust/crates/host/src/machine/mod.rs @@ -1,16 +1,16 @@ mod auth; -mod route_machine; +mod call_machine; use std::future::Future; use std::pin::Pin; -pub use auth::{HostTokenProvider, TokenRoute}; -pub use route_machine::{ExecuteFuture, HostChannel, MachineFault, RouteMachine}; +pub use auth::{HostTokenProvider, TokenProtocol}; +pub use call_machine::{CallMachine, ExecuteFuture, HostChannel, MachineFault}; -use crate::host::{HostOp, HostResult}; -use crate::route::Route; +use crate::host::HostOp; +use crate::protocol::Protocol; -pub enum MachineStep { +pub enum MachineStep { Host(HostOp), Complete(C), } @@ -19,8 +19,8 @@ pub type Step<'a, M> = Pin< Box< dyn Future< Output = Result< - MachineStep<::Route, ::Complete>, - <::Route as Route>::Error, + MachineStep<::Protocol, ::Complete>, + <::Protocol as Protocol>::Error, >, > + Send + 'a, @@ -30,7 +30,10 @@ pub type Step<'a, M> = Pin< pub type Interrupted<'a, M> = Pin< Box< dyn Future< - Output = Result<::Complete, <::Route as Route>::Error>, + Output = Result< + ::Complete, + <::Protocol as Protocol>::Error, + >, > + Send + 'a, >, @@ -51,19 +54,18 @@ impl HostFailure { } /// A resumable call. Core implements it per route; a host drives it. Every suspension -/// point is an op the host performs and answers with a result. +/// point is an op the host performs and answers through the op's own reply before it +/// resumes the call again. pub trait Machine: Send { - type Route: Route; + type Protocol: Protocol; type Complete: Send + 'static; - /// `None` on the first call and whenever the previous step completed without - /// yielding an op; otherwise the result of the op last yielded. - fn resume(&mut self, result: Option>) -> Step<'_, Self>; + fn resume(&mut self) -> Step<'_, Self>; /// The host failed to perform the pending op, or the caller cancelled. The call /// yields no further ops. fn interrupt( &mut self, - failure: HostFailure<::Error>, + failure: HostFailure<::Error>, ) -> Interrupted<'_, Self>; } diff --git a/litellm-rust/crates/host/src/machine/route_machine.rs b/litellm-rust/crates/host/src/machine/route_machine.rs deleted file mode 100644 index 38a0b8bc16a..00000000000 --- a/litellm-rust/crates/host/src/machine/route_machine.rs +++ /dev/null @@ -1,199 +0,0 @@ -//! The one machine every route runs on: it owns the route's provider future, polls it in -//! place, and turns the host operations that future requests into [`Machine`] steps. No -//! task is spawned; dropping the machine drops the in-flight call. - -use std::{future::Future, pin::Pin}; - -use tokio::sync::{mpsc, oneshot}; - -use super::{HostFailure, Interrupted, Machine, MachineStep, Step}; -use crate::{ - event::{MachineEvent, RequestContext, WireRequest}, - host::{Demand, HostOp, HostResult}, - route::Route, -}; - -/// The machine's own failures, distinct from anything the provider call reports. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum MachineFault { - /// The host driver went away while the call was waiting on it. - Abandoned, - /// The host answered out of turn: a result with nothing pending, or nothing when a - /// result was pending. - Protocol(&'static str), - /// The host answered a route operation with the wrong result variant. - Mismatch, -} - -pub type ExecuteFuture = - Pin::Response, ::Error>> + Send>>; - -struct PendingOp { - op: HostOp, - reply: oneshot::Sender>, -} - -/// The provider side of the machine: how the in-flight call reaches its host. -pub struct HostChannel { - ops: mpsc::UnboundedSender>, -} - -impl Clone for HostChannel { - fn clone(&self) -> Self { - Self { - ops: self.ops.clone(), - } - } -} - -impl HostChannel -where - R::Error: From, -{ - async fn invoke(&self, op: HostOp) -> Result, R::Error> { - let (reply, answer) = oneshot::channel(); - self.ops - .send(PendingOp { op, reply }) - .map_err(|_| MachineFault::Abandoned)?; - answer.await.map_err(|_| MachineFault::Abandoned.into()) - } - - pub async fn route(&self, op: R::Op) -> Result { - match self.invoke(HostOp::Route(op)).await? { - HostResult::Route(result) => Ok(result), - _ => Err(MachineFault::Mismatch.into()), - } - } - - pub async fn before_send( - &self, - wire: WireRequest, - context: RequestContext, - ) -> Result { - let op = HostOp::BeforeSend { - wire: Box::new(wire), - context: Box::new(context), - }; - match self.invoke(op).await? { - HostResult::BeforeSend(wire) => Ok(*wire), - _ => Err(MachineFault::Mismatch.into()), - } - } - - pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { - match self.invoke(HostOp::Emit(event)).await? { - HostResult::Emitted => Ok(()), - _ => Err(MachineFault::Mismatch.into()), - } - } - - pub async fn open(&self, head: R::StreamHead) -> Result { - self.demand(HostOp::Open(head)).await - } - - pub async fn deliver(&self, chunk: R::Chunk) -> Result { - self.demand(HostOp::Deliver(chunk)).await - } - - async fn demand(&self, op: HostOp) -> Result { - match self.invoke(op).await? { - HostResult::Demand(demand) => Ok(demand), - _ => Err(MachineFault::Mismatch.into()), - } - } -} - -enum Execution { - Unstarted(Box) -> ExecuteFuture + Send>), - Running(ExecuteFuture), - Done, -} - -pub struct RouteMachine { - execution: Execution, - ops: mpsc::UnboundedReceiver>, - channel: HostChannel, - reply: Option>>, -} - -impl RouteMachine -where - R::Error: From, -{ - pub fn new(execute: impl FnOnce(HostChannel) -> ExecuteFuture + Send + 'static) -> Self { - let (ops_tx, ops) = mpsc::unbounded_channel(); - Self { - execution: Execution::Unstarted(Box::new(execute)), - ops, - channel: HostChannel { ops: ops_tx }, - reply: None, - } - } - - async fn step( - &mut self, - result: Option>, - ) -> Result, R::Error> { - match (self.reply.take(), result) { - (Some(reply), Some(result)) => { - reply - .send(result) - .map_err(|_| MachineFault::Protocol("the call stopped waiting on the host"))?; - } - (None, None) if matches!(self.execution, Execution::Unstarted(_)) => {} - (Some(reply), None) => { - self.reply = Some(reply); - return Err(MachineFault::Protocol("host operation result is required").into()); - } - (None, Some(_)) => { - return Err(MachineFault::Protocol("unexpected host operation result").into()); - } - (None, None) => { - return Err( - MachineFault::Protocol("call cannot be resumed after completion").into(), - ); - } - } - if let Execution::Unstarted(_) = self.execution { - let Execution::Unstarted(start) = - std::mem::replace(&mut self.execution, Execution::Done) - else { - unreachable!() - }; - self.execution = Execution::Running(start(self.channel.clone())); - } - let Execution::Running(future) = &mut self.execution else { - return Err(MachineFault::Protocol("call cannot be resumed after completion").into()); - }; - tokio::select! { - biased; - pending = self.ops.recv() => { - let pending = pending.ok_or(MachineFault::Abandoned)?; - self.reply = Some(pending.reply); - Ok(MachineStep::Host(pending.op)) - } - outcome = future => { - self.execution = Execution::Done; - outcome.map(MachineStep::Complete) - } - } - } -} - -impl Machine for RouteMachine -where - R::Error: From, -{ - type Route = R; - type Complete = R::Response; - - fn resume(&mut self, result: Option>) -> Step<'_, Self> { - Box::pin(self.step(result)) - } - - fn interrupt(&mut self, failure: HostFailure) -> Interrupted<'_, Self> { - self.reply = None; - self.execution = Execution::Done; - Box::pin(async move { Err(failure.into_error()) }) - } -} diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs new file mode 100644 index 00000000000..a7c0f3470b2 --- /dev/null +++ b/litellm-rust/crates/host/src/protocol.rs @@ -0,0 +1,17 @@ +/// One public call surface: what a completed call produces, how it fails, what the host +/// projects the caller's request into, and the protocol-specific operations only its host +/// can perform mid-call (token acquisition, for one). +pub trait Protocol: Send + Sync + 'static { + type Response: Send + 'static; + type Error: Clone + Send + Sync + 'static; + /// The caller's request as the host projects it, answered once before anything else. + type Projection: Send + 'static; + /// Each operation carries the [`Reply`](crate::host::Reply) its answer goes through. + /// A protocol with no operations of its own uses `Infallible`. + type Op: Send + 'static; + /// One piece of a streamed response, handed to the caller as it arrives. A protocol + /// that never streams uses `Infallible`. + type Chunk: Send + 'static; + /// What the call knows once a streamed response starts, before its first chunk. + type StreamHead: Send + 'static; +} diff --git a/litellm-rust/crates/host/src/route.rs b/litellm-rust/crates/host/src/route.rs deleted file mode 100644 index 8ab2b125760..00000000000 --- a/litellm-rust/crates/host/src/route.rs +++ /dev/null @@ -1,14 +0,0 @@ -/// One public call surface: what a completed call produces, how it fails, and the -/// route-specific operations only its host can perform (request projection, file reads, -/// token acquisition). -pub trait Route: Send + Sync + 'static { - type Response: Send + 'static; - type Error: Clone + Send + Sync + 'static; - type Op: Send + 'static; - type OpResult: Send + 'static; - /// One piece of a streamed response, handed to the caller as it arrives. A route - /// that never streams uses `Infallible`. - type Chunk: Send + 'static; - /// What the route knows once a streamed response starts, before its first chunk. - type StreamHead: Send + 'static; -} diff --git a/litellm-rust/crates/host/src/run.rs b/litellm-rust/crates/host/src/run.rs index 6a0c08fba68..baa3b58e058 100644 --- a/litellm-rust/crates/host/src/run.rs +++ b/litellm-rust/crates/host/src/run.rs @@ -1,40 +1,28 @@ use crate::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; -use crate::host::{Host, HostOp, HostResult}; +use crate::host::{Host, HostOp}; use crate::machine::{HostFailure, Machine, MachineStep}; -use crate::route::Route; +use crate::protocol::Protocol; /// Drives a machine to completion against an in-process host and emits exactly one /// terminal event. -pub async fn run(mut machine: M, host: &H) -> Result::Error> +pub async fn run( + mut machine: M, + host: &H, +) -> Result::Error> where M: Machine, - H: Host, + H: Host, { let start_time = epoch_seconds(); let _ = host.emit(&CallEvent::Started { start_time }).await; - let mut result = None; let outcome = loop { - let step = match machine.resume(result.take()).await { + let op = match machine.resume().await { Ok(MachineStep::Complete(complete)) => break Ok(complete), Ok(MachineStep::Host(op)) => op, Err(error) => break Err(error), }; - let answer = match step { - HostOp::Route(op) => host.route(op).await.map(HostResult::Route), - HostOp::BeforeSend { wire, context } => host - .before_send(*wire, &context) - .await - .map(|wire| HostResult::BeforeSend(Box::new(wire))), - HostOp::Emit(event) => host - .emit(&CallEvent::Machine(event)) - .await - .map(|()| HostResult::Emitted), - HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), - HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), - }; - match answer { - Ok(answer) => result = Some(answer), - Err(error) => break machine.interrupt(HostFailure::Error(error)).await, + if let Err(error) = perform(host, op).await { + break machine.interrupt(HostFailure::Error(error)).await; } }; let timing = Timing { @@ -52,44 +40,52 @@ where outcome } +async fn perform>(host: &H, op: HostOp) -> Result<(), R::Error> { + match op { + HostOp::Project(reply) => host + .project() + .await + .map(|projection| reply.send(projection)), + HostOp::Custom(op) => host.custom_op(op).await, + HostOp::BeforeSend { + wire, + context, + reply, + } => host + .before_send(*wire, &context) + .await + .map(|wire| reply.send(wire)), + HostOp::Emit(event, reply) => host + .emit(&CallEvent::Machine(event)) + .await + .map(|()| reply.send(())), + HostOp::Open(head, reply) => host.open(head).await.map(|demand| reply.send(demand)), + HostOp::Deliver(chunk, reply) => host.deliver(chunk).await.map(|demand| reply.send(demand)), + } +} + #[cfg(test)] mod tests { use std::sync::Mutex; use super::*; - use crate::machine::{Interrupted, Step}; + use crate::host::Reply; + use crate::machine::{CallMachine, MachineFault}; struct Unit; - impl Route for Unit { + impl Protocol for Unit { type Response = (); type Error = &'static str; - type Op = &'static str; - type OpResult = (); + type Projection = (); + type Op = (&'static str, Reply<()>); type Chunk = std::convert::Infallible; type StreamHead = std::convert::Infallible; } - struct Scripted { - ops: Vec<&'static str>, - outcome: Result<(), &'static str>, - } - - impl Machine for Scripted { - type Route = Unit; - type Complete = (); - - fn resume(&mut self, _: Option>) -> Step<'_, Self> { - Box::pin(async move { - if !self.ops.is_empty() { - return Ok(MachineStep::Host(HostOp::Route(self.ops.remove(0)))); - } - self.outcome.map(MachineStep::Complete) - }) - } - - fn interrupt(&mut self, failure: HostFailure<&'static str>) -> Interrupted<'_, Self> { - Box::pin(async move { Err(failure.into_error()) }) + impl From for &'static str { + fn from(_: MachineFault) -> Self { + "machine fault" } } @@ -100,12 +96,21 @@ mod tests { } impl Host for Recording { - async fn route(&self, op: &'static str) -> Result<(), &'static str> { - self.seen.lock().unwrap().push(format!("route:{op}")); - match self.fail { - Some(failing) if failing == op => Err("host failed"), - _ => Ok(()), + async fn project(&self) -> Result<(), &'static str> { + self.seen.lock().unwrap().push("project".into()); + Ok(()) + } + + async fn custom_op( + &self, + (op, reply): (&'static str, Reply<()>), + ) -> Result<(), &'static str> { + self.seen.lock().unwrap().push(format!("op:{op}")); + if self.fail == Some(op) { + return Err("host failed"); } + reply.send(()); + Ok(()) } async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { @@ -119,21 +124,29 @@ mod tests { } } - fn scripted(ops: &[&'static str], outcome: Result<(), &'static str>) -> Scripted { - Scripted { - ops: ops.to_vec(), - outcome, - } + fn scripted( + ops: &'static [&'static str], + outcome: Result<(), &'static str>, + ) -> CallMachine { + CallMachine::new(move |host| { + Box::pin(async move { + host.project().await?; + for op in ops { + host.custom_op(|reply| (*op, reply)).await?; + } + outcome + }) + }) } #[tokio::test] async fn forwards_every_op_then_emits_one_succeeded() { let host = Recording::default(); - let outcome = run(scripted(&["project", "send"], Ok(())), &host).await; + let outcome = run(scripted(&["sign", "send"], Ok(())), &host).await; assert_eq!(outcome, Ok(())); assert_eq!( *host.seen.lock().unwrap(), - ["started", "route:project", "route:send", "succeeded"] + ["started", "project", "op:sign", "op:send", "succeeded"] ); } @@ -142,24 +155,32 @@ mod tests { let host = Recording::default(); let outcome = run(scripted(&[], Err("boom")), &host).await; assert_eq!(outcome, Err("boom")); - assert_eq!(*host.seen.lock().unwrap(), ["started", "failed"]); + assert_eq!(*host.seen.lock().unwrap(), ["started", "project", "failed"]); let host = Recording { fail: Some("send"), ..Recording::default() }; - let outcome = run(scripted(&["project", "send", "never"], Ok(())), &host).await; + let outcome = run(scripted(&["sign", "send", "never"], Ok(())), &host).await; assert_eq!(outcome, Err("host failed")); assert_eq!( *host.seen.lock().unwrap(), - ["started", "route:project", "route:send", "failed"] + ["started", "project", "op:sign", "op:send", "failed"] ); } struct StartTimes(Mutex>); impl Host for StartTimes { - async fn route(&self, _: &'static str) -> Result<(), &'static str> { + async fn project(&self) -> Result<(), &'static str> { + Ok(()) + } + + async fn custom_op( + &self, + (_, reply): (&'static str, Reply<()>), + ) -> Result<(), &'static str> { + reply.send(()); Ok(()) } @@ -178,7 +199,7 @@ mod tests { #[tokio::test] async fn started_opens_the_call_at_the_terminal_start_time_and_cannot_fail_it() { let host = StartTimes(Mutex::default()); - assert_eq!(run(scripted(&["project"], Ok(())), &host).await, Ok(())); + assert_eq!(run(scripted(&["send"], Ok(())), &host).await, Ok(())); let times = host.0.lock().unwrap(); assert_eq!(times.len(), 2); assert_eq!(times[0], times[1]); diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index b3df8fc18c8..5d7be8b10df 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -118,7 +118,6 @@ impl From for Error { Self::InvalidRequest(match fault { MachineFault::Abandoned => "OCR host driver was abandoned".into(), MachineFault::Protocol(message) => format!("OCR {message}"), - MachineFault::Mismatch => "invalid OCR host operation result".into(), }) } } diff --git a/litellm-rust/crates/python-bridge/src/logger/machine.rs b/litellm-rust/crates/python-bridge/src/logger/machine.rs index 7234308e67e..54f8f3d4b3f 100644 --- a/litellm-rust/crates/python-bridge/src/logger/machine.rs +++ b/litellm-rust/crates/python-bridge/src/logger/machine.rs @@ -1,9 +1,8 @@ use std::sync::OnceLock; use litellm_host::{ - host::HostResult, machine::{HostFailure, Interrupted, Machine, Step}, - route::Route, + protocol::Protocol, }; use litellm_tracing::Logger; use pyo3::Python; @@ -23,17 +22,17 @@ impl LoggedMachine { } impl Machine for LoggedMachine { - type Route = M::Route; + type Protocol = M::Protocol; type Complete = M::Complete; - fn resume(&mut self, result: Option>) -> Step<'_, Self> { + fn resume(&mut self) -> Step<'_, Self> { let logger = self.logger.get_or_init(|| Python::attach(super::capture)); - Box::pin(logger.instrument(logger.scope(|| self.machine.resume(result)))) + Box::pin(logger.instrument(logger.scope(|| self.machine.resume()))) } fn interrupt( &mut self, - failure: HostFailure<::Error>, + failure: HostFailure<::Error>, ) -> Interrupted<'_, Self> { let logger = self.logger.get_or_init(|| Python::attach(super::capture)); Box::pin(logger.instrument(logger.scope(|| self.machine.interrupt(failure)))) diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 9312d4c187c..b65e7d37023 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -1,29 +1,28 @@ use std::{process::Command, task::Poll}; use litellm_host::{ - host::HostResult, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, - route::Route, + protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; struct DiagnosticMachine; -impl Route for DiagnosticMachine { +impl Protocol for DiagnosticMachine { type Response = (); type Error = String; + type Projection = (); type Op = (); - type OpResult = (); type Chunk = (); type StreamHead = (); } impl Machine for DiagnosticMachine { - type Route = Self; + type Protocol = Self; type Complete = (); - fn resume(&mut self, _: Option>) -> Step<'_, Self> { + fn resume(&mut self) -> Step<'_, Self> { litellm_tracing::warn!("machine started"); Box::pin(async { tokio::task::yield_now().await; @@ -45,7 +44,7 @@ fn machine_warning(py: Python<'_>) -> PyResult> { let mut machine = super::LoggedMachine::new(DiagnosticMachine); let mut future = Box::pin(async move { machine - .resume(None) + .resume() .await .map_err(pyo3::exceptions::PyValueError::new_err)?; machine 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 9d97094aeda..6de4e1320e1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -1,10 +1,12 @@ +use std::convert::Infallible; + use bytes::Bytes; use litellm_core::messages::{ Error, - route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, + route::{Messages, MessagesCall, MessagesOutput}, types::MessagesShaping, }; -use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; +use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; use litellm_http::transport::Error as TransportError; use litellm_types::utils::ProviderSpecificHeaders; use pyo3::{ @@ -80,16 +82,16 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult { /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. -pub(super) struct MessagesRouteHost { +pub(super) struct MessagesPythonHost { request: Py, } -impl MessagesRouteHost { +impl MessagesPythonHost { pub(super) fn new(request: Py) -> Self { Self { request } } - fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + fn projection(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { let request = self.request.bind(py); let argument = |name: &str| -> PyResult>> { Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) @@ -208,22 +210,21 @@ impl MessagesRouteHost { } } -impl RouteHost for MessagesRouteHost { - type Route = Messages; +impl ProtocolHost for MessagesPythonHost { + type Protocol = Messages; type Failure = PyErr; - fn invoke( + fn project( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - op: MessagesOp, - ) -> Result> { - match op { - MessagesOp::ProjectRequest => self - .project(py, arguments) - .map(|call| MessagesOpResult::Request(Box::new(call))) - .map_err(|error| InvokeError::Python(self.map_failure(py, error))), - } + ) -> Result> { + self.projection(py, arguments) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))) + } + + fn invoke(&mut self, _: Python<'_>, op: Infallible) -> Result<(), InvokeError> { + match op {} } fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult> { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index fd474e6b2d4..65040f31684 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,6 +1,6 @@ mod host; -use host::MessagesRouteHost; +use host::MessagesPythonHost; use litellm_callbacks_legacy_python::{ LegacySurface, PassThroughStream, PublicCall, run_legacy_call, }; @@ -45,7 +45,7 @@ fn run_messages( SURFACE, PublicCall::capture(&request, &args, &kwargs)?, crate::logger::LoggedMachine::new(messages_machine(secrets)), - MessagesRouteHost::new(request.unbind()), + MessagesPythonHost::new(request.unbind()), asynchronous, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index a928e62d5b7..e6821241c89 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -1,58 +1,38 @@ use std::path::PathBuf; -use bytes::Bytes; -use litellm_core::ocr::types::{OcrDocumentInput, OcrFileContent}; +use litellm_core::ocr::types::OcrDocumentInput; +use litellm_host_python::{PythonFileReader, py_bytes}; use pyo3::{ - exceptions::{PyTypeError, PyValueError}, - gc::{PyTraverseError, PyVisit}, + exceptions::PyValueError, prelude::*, - pybacked::PyBackedBytes, sync::PyOnceLock, types::{PyBytes, PyString, PyType}, }; -#[derive(Debug)] -pub(super) struct PythonFileReader { - reader: Py, - name: Option, +/// A `type='file'` document as projected: paths and bytes are typed inputs already; a +/// file-like object is a reader the projection consumes once every other field is read. +pub(super) enum FileDocumentInput { + Ready(OcrDocumentInput), + Deferred { + reader: PythonFileReader, + mime_type: Option, + }, } -impl PythonFileReader { - pub(super) fn read(&self, py: Python<'_>) -> PyResult { - let value = self.reader.bind(py).call0()?; - let bytes = if value.is_instance_of::() { - Bytes::from(value.extract::()?) - } else if value.is_instance_of::() { - extract_bytes(&value)? - } else { - return Err(PyTypeError::new_err(format!( - "OCR file read must return bytes or str, got {}", - value.get_type(), - ))); - }; - Ok(OcrFileContent { - bytes, - file_name: self.name.clone(), - }) +impl FileDocumentInput { + pub(super) fn resolve(self, py: Python<'_>) -> PyResult { + match self { + Self::Ready(input) => Ok(input), + Self::Deferred { reader, mime_type } => { + let content = reader.read(py)?; + Ok(OcrDocumentInput::Bytes { + bytes: content.bytes, + file_name: content.file_name, + mime_type, + }) + } + } } - - pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.reader) - } -} - -fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult { - if value.is_exact_instance_of::() { - return Ok(Bytes::from_owner(value.extract::()?)); - } - Ok(Bytes::copy_from_slice( - value.extract::()?.as_ref(), - )) -} - -pub(super) struct FileDocumentInput { - pub input: OcrDocumentInput, - pub reader: Option, } impl FromPyObject<'_, '_> for FileDocumentInput { @@ -87,51 +67,31 @@ impl FromPyObject<'_, '_> for FileDocumentInput { } static PATH_LIKE: PyOnceLock> = PyOnceLock::new(); if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? { - return Ok(Self { - input: OcrDocumentInput::Path { - path: file.extract::()?, - mime_type, - }, - reader: None, - }); + return Ok(Self::Ready(OcrDocumentInput::Path { + path: file.extract::()?, + mime_type, + })); } if file.is_instance_of::() { - return Ok(Self { - input: OcrDocumentInput::Bytes { - bytes: extract_bytes(&file)?, - file_name: None, - mime_type, - }, - reader: None, - }); + return Ok(Self::Ready(OcrDocumentInput::Bytes { + bytes: py_bytes(&file)?, + file_name: None, + mime_type, + })); } - let reader = file - .getattr_opt("read")? - .filter(|value| value.is_callable()); - let Some(reader) = reader else { - return Err(PyValueError::new_err(format!( + match PythonFileReader::from_file_like(&file)? { + Some(reader) => Ok(Self::Deferred { reader, mime_type }), + None => Err(PyValueError::new_err(format!( "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", file.get_type(), - ))); - }; - let name = file - .getattr_opt("name")? - .filter(|value| !value.is_none()) - .map(|value| value.extract::()) - .transpose()?; - Ok(Self { - input: OcrDocumentInput::HostReader { mime_type }, - reader: Some(PythonFileReader { - reader: reader.unbind(), - name, - }), - }) + ))), + } } } #[cfg(test)] mod tests { - use pyo3::types::PyDict; + use pyo3::{exceptions::PyTypeError, types::PyDict}; use super::*; @@ -141,6 +101,13 @@ mod tests { locals } + fn ready(input: FileDocumentInput) -> OcrDocumentInput { + match input { + FileDocumentInput::Ready(input) => input, + FileDocumentInput::Deferred { .. } => panic!("expected a ready document"), + } + } + #[test] fn extraction_validates_required_file_and_optional_mime_type() { Python::initialize(); @@ -167,13 +134,19 @@ mod tests { .unwrap(); assert!(error.is_instance_of::(py)); assert!(error.to_string().contains("bare str")); + let error = py + .eval(c"{'file': object()}", None, None) + .unwrap() + .extract::() + .err() + .unwrap(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("Unsupported file input type")); let document = py .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None) .unwrap(); - let input: FileDocumentInput = document.extract().unwrap(); - assert!(input.reader.is_none()); assert_eq!( - input.input, + ready(document.extract().unwrap()), OcrDocumentInput::Bytes { bytes: b"abc".as_slice().into(), file_name: None, @@ -199,7 +172,7 @@ class Reader: return b'abc' reader = Reader() document = {'file': reader, 'mime_type': 7} -reader_document = {'file': reader} +reader_document = {'file': reader, 'mime_type': 'application/pdf'} path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}", ); let document = locals.get_item("document").unwrap().unwrap(); @@ -208,10 +181,6 @@ path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_typ let document = locals.get_item("reader_document").unwrap().unwrap(); let input: FileDocumentInput = document.extract().unwrap(); - assert_eq!( - input.input, - OcrDocumentInput::HostReader { mime_type: None } - ); let reads = || { locals .get_item("reader") @@ -223,21 +192,20 @@ path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_typ .unwrap() }; assert_eq!(reads(), 0); - let content = input.reader.unwrap().read(py).unwrap(); + let resolved = input.resolve(py).unwrap(); assert_eq!(reads(), 1); assert_eq!( - content, - OcrFileContent { + resolved, + OcrDocumentInput::Bytes { bytes: b"abc".as_slice().into(), file_name: Some("scan.png".into()), + mime_type: Some("application/pdf".into()), } ); let document = locals.get_item("path_document").unwrap().unwrap(); - let input: FileDocumentInput = document.extract().unwrap(); - assert!(input.reader.is_none()); assert_eq!( - input.input, + ready(document.extract().unwrap()), OcrDocumentInput::Path { path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"), mime_type: Some("image/png".into()), @@ -245,97 +213,4 @@ path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_typ ); }); } - - #[test] - fn reader_results_are_normalized_and_exceptions_keep_their_identity() { - Python::initialize(); - Python::attach(|py| { - let locals = eval( - py, - c"failure = KeyError('reader failed') -class Raising: - def read(self): - raise failure -class Text: - def read(self): - return 'héllo' -class Wrong: - def read(self): - return 7 -raising = {'file': Raising()} -text = {'file': Text()} -wrong = {'file': Wrong()}", - ); - let reader = |name: &str| { - locals - .get_item(name) - .unwrap() - .unwrap() - .extract::() - .unwrap() - .reader - .unwrap() - }; - let error = reader("raising").read(py).unwrap_err(); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - assert_eq!( - reader("text").read(py).unwrap().bytes.as_ref(), - "héllo".as_bytes() - ); - let error = reader("wrong").read(py).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("bytes or str")); - }); - } - - #[rstest::rstest] - #[case::read("read")] - #[case::name("name")] - fn reader_attribute_failures_keep_their_identity(#[case] attribute: &str) { - Python::initialize(); - Python::attach(|py| { - let locals = eval( - py, - c"failure = LookupError('file property failed') -class File: - def __getattribute__(self, name): - if name == attribute: - raise failure - return super().__getattribute__(name) - name = 'scan.pdf' - def read(self): - return b'abc' -document = {'file': File()}", - ); - locals.set_item("attribute", attribute).unwrap(); - let error = locals - .get_item("document") - .unwrap() - .unwrap() - .extract::() - .err() - .unwrap(); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - }); - } - - #[test] - fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() { - Python::initialize(); - let (bytes, pointer) = Python::attach(|py| { - let value = PyBytes::new(py, b"document bytes"); - let pointer = value.as_bytes().as_ptr() as usize; - (extract_bytes(value.as_any()).unwrap(), pointer) - }); - assert_eq!(bytes.as_ptr() as usize, pointer); - assert_eq!(bytes.as_ref(), b"document bytes"); - } } 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 8bf99cd355f..5a3806e61e3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,6 +1,6 @@ use litellm_auth::ResolvedCredential; -use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{InvokeError, RouteHost, missing_state, to_py}; +use litellm_core::ocr::route::{Ocr, OcrOp, OcrProjection}; +use litellm_host_python::{InvokeError, ProtocolHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -20,14 +20,15 @@ enum OcrHostData { Released, } -/// The Python side of the OCR route: projects the prepared arguments, reads file-like -/// documents, acquires Azure AD tokens, and builds the public response and exception. -pub(super) struct OcrRouteHost { +/// The Python side of the OCR route: projects the prepared arguments (reading a file-like +/// document as it goes), acquires Azure AD tokens, and builds the public response and +/// exception. +pub(super) struct OcrPythonHost { request: Py, data: OcrHostData, } -impl OcrRouteHost { +impl OcrPythonHost { pub(super) fn new(request: Py) -> Self { Self { request, @@ -42,14 +43,6 @@ impl OcrRouteHost { } } - fn read_document(&self, py: Python<'_>) -> PyResult { - self.handles()? - .reader - .as_ref() - .ok_or_else(missing_state)? - .read(py) - } - fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { self.handles()? .azure_ad_token_provider @@ -58,30 +51,21 @@ impl OcrRouteHost { .acquire(py) } - fn answer( + fn projection( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - op: OcrOp, - ) -> PyResult { - match op { - OcrOp::ProjectRequest => { - let OcrHostData::Unprojected = self.data else { - return Err(missing_state()); - }; - let (request, handles) = project_request(self.request.bind(py), arguments)?; - let caller_token = handles.azure_ad_token_provider.is_some(); - self.data = OcrHostData::Projected(Box::new(handles)); - Ok(OcrOpResult::Request { - request: Box::new(request), - caller_token, - }) - } - OcrOp::ReadDocument => self.read_document(py).map(OcrOpResult::Document), - OcrOp::AcquireAzureAdToken => self - .acquire_azure_ad_token(py) - .map(OcrOpResult::AzureAdToken), - } + ) -> PyResult { + let OcrHostData::Unprojected = self.data else { + return Err(missing_state()); + }; + let (request, handles) = project_request(self.request.bind(py), arguments)?; + let caller_token = handles.azure_ad_token_provider.is_some(); + self.data = OcrHostData::Projected(Box::new(handles)); + Ok(OcrProjection { + request, + caller_token, + }) } fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { @@ -104,20 +88,28 @@ impl OcrRouteHost { } } -impl RouteHost for OcrRouteHost { - type Route = Ocr; +impl ProtocolHost for OcrPythonHost { + type Protocol = Ocr; type Failure = PyErr; - fn invoke( + fn project( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - op: OcrOp, - ) -> Result> { - self.answer(py, arguments, op) + ) -> Result> { + self.projection(py, arguments) .map_err(|error| InvokeError::Python(self.map_failure(py, error))) } + fn invoke(&mut self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError> { + match op { + OcrOp::AcquireAzureAdToken(reply) => self + .acquire_azure_ad_token(py) + .map(|token| reply.send(token)) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))), + } + } + fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { py.import("litellm.rust_bridge.ocr.route_host")? .getattr("response")? @@ -148,13 +140,10 @@ impl RouteHost for OcrRouteHost { fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { visit.call(&self.request)?; - if let OcrHostData::Projected(handles) = &self.data { - if let Some(reader) = &handles.reader { - reader.traverse(visit)?; - } - if let Some(provider) = &handles.azure_ad_token_provider { - provider.traverse(visit)?; - } + if let OcrHostData::Projected(handles) = &self.data + && let Some(provider) = &handles.azure_ad_token_provider + { + provider.traverse(visit)?; } Ok(()) } @@ -205,20 +194,13 @@ del provider .unwrap() .cast_into::() .unwrap(); - let mut host = OcrRouteHost::new(py.None()); - let projected = host.invoke(py, &kwargs, OcrOp::ProjectRequest).unwrap(); - assert!(matches!( - projected, - OcrOpResult::Request { - caller_token: true, - .. - } - )); + let mut host = OcrPythonHost::new(py.None()); + assert!(host.project(py, &kwargs).unwrap().caller_token); locals.del_item("kwargs").unwrap(); drop(kwargs); + let (reply, _) = litellm_host::host::reply(); assert_eq!( - host.invoke(py, &PyDict::new(py), OcrOp::AcquireAzureAdToken) - .is_ok(), + host.invoke(py, OcrOp::AcquireAzureAdToken(reply)).is_ok(), succeeds ); let alive = || { diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index b54316b258b..a4f2bf851d7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -5,7 +5,7 @@ mod project; use std::sync::LazyLock; -use host::OcrRouteHost; +use host::OcrPythonHost; use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::{provider_config, route::ocr_machine}; @@ -69,7 +69,7 @@ fn run_ocr( if asynchronous { ASYNC_SURFACE } else { SURFACE }, PublicCall::capture(&request, &args, &kwargs)?, crate::logger::LoggedMachine::new(ocr_machine(client)), - OcrRouteHost::new(request.unbind()), + OcrPythonHost::new(request.unbind()), asynchronous, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 697b935a1d4..be43a1b7711 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -8,19 +8,15 @@ use litellm_llms::base_llm::ocr::error::Error; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use serde_json::{Map, Value}; -use super::{ - document::{FileDocumentInput, PythonFileReader}, - errors::to_pyerr as ocr_error_to_pyerr, -}; +use super::{document::FileDocumentInput, errors::to_pyerr as ocr_error_to_pyerr}; use crate::{ credentials::{self, CallerTokenProvider}, marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}, }; -/// What the host keeps after projection: the caller's callables that answer the document -/// read and token operations, and the provider name the failure mapping reports. +/// What the host keeps after projection: the caller's token callable that answers the +/// token operation, and the provider name the failure mapping reports. pub(super) struct OcrHostHandles { - pub reader: Option, pub azure_ad_token_provider: Option, pub provider: &'static str, } @@ -104,13 +100,11 @@ impl ProjectedDocument { Ok(Self::File(document.extract()?)) } - fn into_parts(self) -> PyResult<(OcrDocumentInput, Option)> { + /// Reads a file-like document now, so it runs after every other argument was read. + fn resolve(self, py: Python<'_>) -> PyResult { match self { - Self::File(FileDocumentInput { input, reader }) => Ok((input, reader)), - Self::Other(wire) => Ok(( - decode_document(wire).map_err(ocr_error_to_pyerr)?.into(), - None, - )), + Self::File(file) => file.resolve(py), + Self::Other(wire) => Ok(decode_document(wire).map_err(ocr_error_to_pyerr)?.into()), } } } @@ -136,24 +130,25 @@ pub(super) fn project_request( .chain(["api_key", "api_base", "extra_headers"]), )?; let azure_ad_token_provider = credentials::azure_ad_token_provider(kwargs)?; - let (document, reader) = document.into_parts()?; + let api_base = arguments.api_base()?; + let extra_headers = arguments.extra_headers()?; + let timeout_seconds = arguments.timeout_seconds()?; let wire = OcrWireRequest { model, - document, + document: document.resolve(request.py())?, api_key, - api_base: arguments.api_base()?, + api_base, custom_llm_provider, - extra_headers: arguments.extra_headers()?, + extra_headers, optional_params, input_sources, - timeout_seconds: arguments.timeout_seconds()?, + timeout_seconds, }; let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?; let provider = request.provider_name(); Ok(( request, OcrHostHandles { - reader, azure_ad_token_provider, provider, }, @@ -180,10 +175,8 @@ mod tests { OcrArguments { request, kwargs } } - fn project_document( - document: &Bound<'_, PyAny>, - ) -> PyResult<(OcrDocumentInput, Option)> { - ProjectedDocument::project(document)?.into_parts() + fn project_document(document: &Bound<'_, PyAny>) -> PyResult { + ProjectedDocument::project(document)?.resolve(document.py()) } fn url_document(url: &str) -> OcrDocumentInput { @@ -342,8 +335,11 @@ kwargs = {} }); } + /// A reader that rewrites the request while it runs shows which arguments projection + /// read before it and which after: every other argument is read first, and the read + /// happens exactly once. #[test] - fn document_readers_are_not_consumed_during_projection() { + fn document_readers_are_read_once_after_every_other_argument() { Python::initialize(); Python::attach(|py| { stub_timeout_conversion(py); @@ -351,17 +347,24 @@ kwargs = {} py, c" class Request: - api_base = 'original' + model = 'mistral/mistral-ocr-latest' + custom_llm_provider = None + api_key = None + api_base = 'https://original.example.com' + extra_headers = {'x-source': 'original'} timeout = 1 @property def document(self): return document class Reader: + reads = 0 def read(self): - Request.api_base = 'mutated' + Reader.reads += 1 + Request.api_base = 'https://mutated.example.com' + Request.extra_headers = {'x-source': 'mutated'} Request.timeout = 9 return b'abc' -document = {'type': 'file', 'file': Reader()} +document = {'type': 'file', 'file': Reader(), 'mime_type': 'application/pdf'} request = Request() kwargs = {} ", @@ -373,15 +376,38 @@ kwargs = {} .unwrap() .cast_into::() .unwrap(); - let arguments = arguments(&request, &kwargs); - let document = arguments.document().unwrap(); - let (input, reader) = project_document(&document).unwrap(); - assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None }); - assert_eq!(arguments.api_base().unwrap().as_deref(), Some("original")); - assert_eq!(arguments.timeout_seconds().unwrap(), Some(1.0)); - reader.unwrap().read(py).unwrap(); - assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated")); - assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0)); + let (projected, _) = project_request(&request, &kwargs).unwrap(); + assert_eq!( + py.eval(c"Reader.reads", Some(&locals), Some(&locals)) + .unwrap() + .extract::() + .unwrap(), + 1 + ); + assert_eq!( + projected.document, + OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + } + ); + assert_eq!( + projected + .credentials + .api_base + .as_ref() + .map(|base| base.value().as_str()), + Some("https://original.example.com") + ); + assert_eq!( + projected.transport.extra_headers, + [("x-source".to_string(), "original".to_string())] + ); + assert_eq!( + projected.transport.timeout, + Some(std::time::Duration::from_secs(1)) + ); }); } @@ -396,16 +422,14 @@ kwargs = {} None, ) .unwrap(); - let (input, reader) = project_document(&file).unwrap(); assert_eq!( - input, + project_document(&file).unwrap(), OcrDocumentInput::Bytes { bytes: b"%PDF-1.4".as_slice().into(), file_name: None, mime_type: Some("application/pdf".into()), } ); - assert!(reader.is_none()); let original = py .eval( @@ -414,8 +438,10 @@ kwargs = {} None, ) .unwrap(); - let (input, _) = project_document(&original).unwrap(); - assert_eq!(input, url_document("https://example.com/a.pdf")); + assert_eq!( + project_document(&original).unwrap(), + url_document("https://example.com/a.pdf") + ); }); } @@ -617,7 +643,7 @@ document = Document() ", ); let document = locals.get_item("document").unwrap().unwrap(); - let (input, _) = project_document(&document).unwrap(); + let input = project_document(&document).unwrap(); assert!(matches!(input, OcrDocumentInput::Bytes { .. })); let reads: Vec = document.getattr("reads").unwrap().extract().unwrap(); assert_eq!(reads, ["type", "mime_type", "file"]);