scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

refactor(rust): extract the host coroutine into its own crate (#43129)

Browse source 
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 <yujong@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-25 04:20:06 +00:00 • committed by GitHub
parent 0d47347ad7
commit 1f8997398e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 1660 additions and 1139 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/AGENTS.md
View file

@ -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/<name>/mod.rs` or `tests/<subject>/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

11
litellm-rust/Cargo.lock generated
View file

@ -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",

1
litellm-rust/Cargo.toml
View file

@ -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" }

14
litellm-rust/crates/callbacks-legacy-python/src/call.rs
View file

@ -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<H, M>(
py: Python<'_>,
surface: LegacySurface,
call: PublicCall,
machine: M,
route: H,
host: H,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: RouteHost + 'static,
M: Machine<Route = H::Route, Complete = <H::Route as Route>::Response> + 'static,
H: ProtocolHost + 'static,
M: Machine<Protocol = H::Protocol, Complete = <H::Protocol as Protocol>::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,

49
litellm-rust/crates/core/src/messages/route.rs
View file

@ -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<MessagesCall>),
}
/// 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<MachineFault> 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<Messages>;
pub type MessagesMachine = RouteMachine<Messages>;
pub type MessagesMachine = CallMachine<Messages>;
/// 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<Messages> for LocalMessagesHost {
async fn route(&self, op: MessagesOp) -> Result<MessagesOpResult, Error> {
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<MessagesCall, Error> {
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<dyn SecretSource>) -> 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<dyn SecretSource>,
) -> Result<MessagesOutput, Error> {
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?;

6
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -23,9 +23,6 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
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]

343
litellm-rust/crates/core/src/ocr/route.rs
View file

@ -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<ResolvedCredential>),
}
pub enum OcrOpResult {
Request {
request: Box<LiteLLMOcrRequest<OcrDocumentInput>>,
caller_token: bool,
},
Document(OcrFileContent),
AzureAdToken(ResolvedCredential),
/// The caller's request as the host projects it.
pub struct OcrProjection {
pub request: LiteLLMOcrRequest<OcrDocumentInput>,
/// 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<ResolvedCredential> {
match result {
OcrOpResult::AzureAdToken(credential) => Some(credential),
_ => None,
}
impl TokenProtocol for Ocr {
fn acquire_token_op(reply: Reply<ResolvedCredential>) -> OcrOp {
OcrOp::AcquireAzureAdToken(reply)
}
}
pub type OcrHost = HostChannel<Ocr>;
pub type OcrMachine = RouteMachine<Ocr>;
pub type OcrMachine = CallMachine<Ocr>;
/// 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<LiteLLMOcrResponse, Error> {
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<OcrDocumentInput>,
host: &OcrHost,
) -> Result<ResolvedOcrRequest, Error> {
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<dyn Fn() -> Result<OcrFileContent, Error> + Send + Sync>;
type BeforeSend =
Box<dyn Fn(WireRequest, &RequestContext) -> Result<WireRequest, Error> + Send + Sync>;
type Observer = Box<dyn Fn(&CallEvent) + Send + Sync>;
@ -116,7 +86,6 @@ type Observer = Box<dyn Fn(&CallEvent) + Send + Sync>;
/// projection, and the optional observer sees and may rewrite the wire request.
pub struct LocalOcrHost {
request: Mutex<Option<LiteLLMOcrRequest<OcrDocumentInput>>>,
reader: Option<Reader>,
before_send: Option<BeforeSend>,
observer: Option<Observer>,
}
@ -125,22 +94,11 @@ impl LocalOcrHost {
pub fn new(request: LiteLLMOcrRequest<OcrDocumentInput>) -> Self {
Self {
request: Mutex::new(Some(request)),
reader: None,
before_send: None,
observer: None,
}
}
pub fn with_reader(
self,
reader: impl Fn() -> Result<OcrFileContent, Error> + Send + Sync + 'static,
) -> Self {
Self {
reader: Some(Box::new(reader)),
..self
}
}
pub fn with_before_send(
self,
before_send: impl Fn(WireRequest, &RequestContext) -> Result<WireRequest, Error>
@ -163,25 +121,21 @@ impl LocalOcrHost {
}
impl litellm_host::host::Host<Ocr> for LocalOcrHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, Error> {
async fn project(&self) -> Result<OcrProjection, Error> {
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<crate::ocr::types::OcrDocumentInput>,
content: Result<crate::ocr::types::OcrFileContent, OcrError>,
) -> (Result<LiteLLMOcrResponse, OcrError>, 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<crate::ocr::route::Ocr> for CallerTokenHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, OcrError> {
async fn project(&self) -> Result<OcrProjection, OcrError> {
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"),
});
}
}
}
}

9
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -25,9 +25,6 @@ pub enum OcrDocumentInput {
file_name: Option<String>,
mime_type: Option<String>,
},
HostReader {
mime_type: Option<String>,
},
}
impl From<OcrDocument> for OcrDocumentInput {
@ -45,12 +42,6 @@ impl From<PathBuf> for OcrDocumentInput {
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OcrFileContent {
pub bytes: Bytes,
pub file_name: Option<String>,
}
/// 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`.

31
litellm-rust/crates/coroutine/AGENTS.md Normal file
View file

@ -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

15
litellm-rust/crates/coroutine/Cargo.toml Normal file
View file

@ -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"] }

42
litellm-rust/crates/coroutine/src/co.rs Normal file
View file

@ -0,0 +1,42 @@
use std::sync::Weak;
use tokio::sync::mpsc;
use crate::{Abandoned, Reply, reply};
pub(crate) struct Request<Y> {
pub(crate) value: Y,
pub(crate) outstanding: Weak<()>,
}
/// The body's handle for yielding, `genawaiter`'s `Co`.
pub struct Co<Y> {
yields: mpsc::UnboundedSender<Request<Y>>,
}
impl<Y> Clone for Co<Y> {
fn clone(&self) -> Self {
Self {
yields: self.yields.clone(),
}
}
}
impl<Y> Co<Y> {
pub(crate) fn new(yields: mpsc::UnboundedSender<Request<Y>>) -> Self {
Self { yields }
}
/// Yields the value `ask` builds around a fresh [`Reply`] and waits for its answer.
pub async fn yield_<A>(&self, ask: impl FnOnce(Reply<A>) -> Y) -> Result<A, Abandoned> {
let (reply, answer) = reply();
let outstanding = reply.outstanding();
self.yields
.send(Request {
value: ask(reply),
outstanding,
})
.map_err(|_| Abandoned)?;
answer.await
}
}

94
litellm-rust/crates/coroutine/src/coroutine.rs Normal file
View file

@ -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<Y, C> {
Yielded(Y),
Complete(C),
}
type Body<C> = Pin<Box<dyn Future<Output = C> + Send>>;
enum Step<Y, C> {
Yielded(Request<Y>),
Complete(C),
}
fn queued<Y>(
yields: &mut mpsc::UnboundedReceiver<Request<Y>>,
context: &mut Context<'_>,
) -> Option<Request<Y>> {
match yields.poll_recv(context) {
Poll::Ready(request) => request,
Poll::Pending => None,
}
}
pub struct Coroutine<Y, C> {
body: Option<Body<C>>,
yields: mpsc::UnboundedReceiver<Request<Y>>,
outstanding: Weak<()>,
}
impl<Y, C> Coroutine<Y, C> {
/// Builds the body from `producer`. Nothing runs until the first `resume`.
pub fn new<F>(producer: impl FnOnce(Co<Y>) -> F) -> Self
where
F: Future<Output = C> + 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<CoroutineState<Y, C>, 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() {}
}
}

14
litellm-rust/crates/coroutine/src/error.rs Normal file
View file

@ -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;

12
litellm-rust/crates/coroutine/src/lib.rs Normal file
View file

@ -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};

60
litellm-rust/crates/coroutine/src/reply.rs Normal file
View file

@ -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<A> {
slot: oneshot::Sender<A>,
outstanding: Arc<()>,
}
impl<A> Reply<A> {
/// 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<A> fmt::Debug for Reply<A> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Reply")
}
}
/// The waiting end of a [`Reply`].
pub struct Answer<A> {
slot: oneshot::Receiver<A>,
}
impl<A> Future for Answer<A> {
type Output = Result<A, Abandoned>;
fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
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<A>() -> (Reply<A>, Answer<A>) {
let (slot, answer) = oneshot::channel();
let reply = Reply {
slot,
outstanding: Arc::new(()),
};
(reply, Answer { slot: answer })
}

256
litellm-rust/crates/coroutine/tests/coroutine.rs Normal file
View file

@ -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<u32>),
}
type Test<C> = Coroutine<Ask, C>;
fn yielded<C>(state: Result<CoroutineState<Ask, C>, 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<C>(state: Result<CoroutineState<Ask, C>, 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<u32> {
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<Result<&'static str, Abandoned>> {
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<String> = 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<u32> = 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<u32>| 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<u32> = 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<u32> = 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<Mutex<bool>>);
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<Mutex<Option<Co<Ask>>>> = 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::<u32>();
if send {
reply.send(5);
} else {
drop(reply);
}
assert_eq!(answer.await, if send { Ok(5) } else { Err(Abandoned) });
}

4
litellm-rust/crates/host-python/AGENTS.md
View file

@ -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

1
litellm-rust/crates/host-python/Cargo.toml
View file

@ -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

39
litellm-rust/crates/host-python/src/adapter.rs
View file

@ -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<E> From<PyErr> for InvokeError<E> {
}
}
/// 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<Error: std::fmt::Display>;
pub trait ProtocolHost: Send + Sync {
type Protocol: Protocol<Error: std::fmt::Display>;
/// The public exception a native failure maps to, kept as a value until the driver
/// raises it.
type Failure: Into<PyErr>;
/// `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: <Self::Route as Route>::Op,
) -> Result<<Self::Route as Route>::OpResult, InvokeError<<Self::Route as Route>::Error>>;
) -> Result<
<Self::Protocol as Protocol>::Projection,
InvokeError<<Self::Protocol as Protocol>::Error>,
>;
/// Answers `op` through its reply.
fn invoke(
&mut self,
py: Python<'_>,
op: <Self::Protocol as Protocol>::Op,
) -> Result<(), InvokeError<<Self::Protocol as Protocol>::Error>>;
fn complete(
&mut self,
py: Python<'_>,
response: <Self::Route as Route>::Response,
response: <Self::Protocol as Protocol>::Response,
) -> PyResult<Py<PyAny>>;
/// One streamed chunk as the caller receives it.
fn chunk(
&mut self,
py: Python<'_>,
chunk: <Self::Route as Route>::Chunk,
chunk: <Self::Protocol as Protocol>::Chunk,
) -> PyResult<Py<PyAny>>;
fn classify(
&self,
py: Python<'_>,
error: <Self::Route as Route>::Error,
error: <Self::Protocol as Protocol>::Error,
) -> PyResult<Self::Failure>;
fn host_error(error: &PyErr) -> <Self::Route as Route>::Error;
fn host_error(error: &PyErr) -> <Self::Protocol as Protocol>::Error;
fn close(&mut self, py: Python<'_>);

401
litellm-rust/crates/host-python/src/driver.rs
View file

@ -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<H> = <H as RouteHost>::Route;
type ErrorOf<H> = <RouteOf<H> as Route>::Error;
type ResponseOf<H> = <RouteOf<H> as Route>::Response;
type NativeStep<H> = MachineStep<RouteOf<H>, ResponseOf<H>>;
type ProtocolOf<H> = <H as ProtocolHost>::Protocol;
type ErrorOf<H> = <ProtocolOf<H> as Protocol>::Error;
type ResponseOf<H> = <ProtocolOf<H> as Protocol>::Response;
type NativeStep<H> = MachineStep<ProtocolOf<H>, ResponseOf<H>>;
type NativeResult<H> = Result<NativeStep<H>, ErrorOf<H>>;
type NativeResume<H> = Option<Result<HostResult<RouteOf<H>>, HostFailure<ErrorOf<H>>>>;
type Interruption<H> = Option<HostFailure<ErrorOf<H>>>;
type MachineResult<M> = Result<
MachineStep<<M as Machine>::Route, <M as Machine>::Complete>,
<<M as Machine>::Route as Route>::Error,
MachineStep<<M as Machine>::Protocol, <M as Machine>::Complete>,
<<M as Machine>::Protocol as Protocol>::Error,
>;
struct MachineState<M: Machine> {
@ -44,12 +45,11 @@ enum Stage {
Failed(Py<PyBaseException>),
}
#[derive(Clone, Copy)]
enum Expect {
Started,
Arguments,
Wire,
Emitted,
Wire(Reply<WireRequest>),
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<Demand>),
}
enum Next<H: RouteHost> {
/// 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<E>(answer: Result<(), InvokeError<E>>) -> PyResult<Result<(), E>> {
match answer {
Ok(()) => Ok(Ok(())),
Err(InvokeError::Native(error)) => Ok(Err(error)),
Err(InvokeError::Python(error)) => Err(error),
}
}
enum Next<H: ProtocolHost> {
Return(ExecutionStep),
Continue(HostStep<NativeResult<H>, Py<PyAny>>),
}
struct PythonDriver<H, M>
where
H: RouteHost,
M: Machine<Route = H::Route, Complete = ResponseOf<H>> + 'static,
H: ProtocolHost,
M: Machine<Protocol = H::Protocol, Complete = ResponseOf<H>> + 'static,
{
route: H,
host: H,
adapter: Box<dyn PythonLifecycle>,
machine: Option<Arc<Mutex<MachineState<M>>>>,
arguments: Option<Py<PyDict>>,
@ -89,17 +99,17 @@ where
pub fn run_call<H, M>(
py: Python<'_>,
machine: M,
route: H,
host: H,
adapter: Box<dyn PythonLifecycle>,
arguments: Py<PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: RouteHost + 'static,
M: Machine<Route = H::Route, Complete = ResponseOf<H>> + 'static,
H: ProtocolHost + 'static,
M: Machine<Protocol = H::Protocol, Complete = ResponseOf<H>> + '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<H, M> PythonDriver<H, M>
where
H: RouteHost,
M: Machine<Route = H::Route, Complete = ResponseOf<H>> + 'static,
H: ProtocolHost,
M: Machine<Protocol = H::Protocol, Complete = ResponseOf<H>> + '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<ExecutionStep> {
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<H>,
interruption: Interruption<H>,
) -> PyResult<ExecutionStep> {
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<ExecutionStep> {
fn opened(&mut self, py: Python<'_>, reply: Reply<Demand>) -> PyResult<ExecutionStep> {
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: <RouteOf<H> as Route>::Chunk,
chunk: <ProtocolOf<H> as Protocol>::Chunk,
reply: Reply<Demand>,
) -> PyResult<ExecutionStep> {
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<H>,
interruption: Interruption<H>,
) -> PyResult<HostStep<NativeResult<H>, Py<PyAny>>> {
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<H>) -> PyResult<ExecutionStep> {
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<H>) -> 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<H, M> ExecutionBody for PythonDriver<H, M>
where
H: RouteHost,
M: Machine<Route = H::Route, Complete = ResponseOf<H>> + 'static,
H: ProtocolHost,
M: Machine<Protocol = H::Protocol, Complete = ResponseOf<H>> + 'static,
{
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
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<H, M> Drop for PythonDriver<H, M>
where
H: RouteHost,
M: Machine<Route = H::Route, Complete = ResponseOf<H>> + 'static,
H: ProtocolHost,
M: Machine<Protocol = H::Protocol, Complete = ResponseOf<H>> + '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<MachineFault> 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<HostOp<Synthetic>>,
outcome: Option<Result<String, Error>>,
answers: Vec<String>,
struct Synthetic;
impl Protocol for Synthetic {
type Response = String;
type Error = Error;
type Projection = String;
type Op = (&'static str, Reply<String>);
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<HostResult<Synthetic>>) -> 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<Error>) -> Interrupted<'_, Self> {
self.ops.clear();
self.outcome = None;
Box::pin(async move { Err(failure.into_error()) })
}
}
#[derive(Default)]
struct Log(Arc<Mutex<Vec<String>>>);
@ -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<String, InvokeError<Error>> {
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<String, InvokeError<Error>> {
self.log.push("project");
self.answer(|| format!("project:{}", arguments.len()))
}
fn invoke(
&mut self,
_: Python<'_>,
arguments: &Bound<'_, PyDict>,
op: &'static str,
) -> Result<String, InvokeError<Error>> {
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<String>),
) -> Result<(), InvokeError<Error>> {
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<Py<PyAny>> {
@ -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<Synthetic>,
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<Synthetic>,
host: SyntheticHost,
script: AdapterScript,
asynchronous: bool,
) -> (PyResult<Py<PyAny>>, Vec<String>) {
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<Synthetic> {
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::<String>(py).unwrap(), "done");
assert_eq!(
result.unwrap().extract::<String>(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<Synthetic> {
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::<String>(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<String, InvokeError<Error>> {
self.0.push("route");
self.0.push("project");
Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
}
fn invoke(
&mut self,
_: Python<'_>,
_: (&'static str, Reply<String>),
) -> Result<(), InvokeError<Error>> {
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::<pyo3::exceptions::PyException>(py));
assert_eq!(
log.entries(),
["started", "begin", "route", "adapter.close"]
["started", "begin", "project", "adapter.close"]
);
});
}

241
litellm-rust/crates/host-python/src/file_reader.rs Normal file
View file

@ -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<String>,
}
#[derive(Debug)]
pub struct PythonFileReader {
reader: Py<PyAny>,
name: Option<String>,
}
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<Option<Self>> {
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::<String>())
.transpose()?;
Ok(Some(Self {
reader: reader.unbind(),
name,
}))
}
pub fn read(&self, py: Python<'_>) -> PyResult<FileContent> {
let value = self.reader.bind(py).call0()?;
let bytes = if value.is_instance_of::<PyString>() {
Bytes::from(value.extract::<String>()?)
} else if value.is_instance_of::<PyBytes>() {
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<Bytes> {
if value.is_exact_instance_of::<PyBytes>() {
return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
}
Ok(Bytes::copy_from_slice(
value.extract::<PyBackedBytes>()?.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::<usize>()
.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::<PyTypeError>(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");
}
}

6
litellm-rust/crates/host-python/src/lib.rs
View file

@ -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};

1
litellm-rust/crates/host/Cargo.toml
View file

@ -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"] }

38
litellm-rust/crates/host/src/host.rs
View file

@ -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<R: Route> {
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<R: Protocol> {
/// The first op of every call: the caller's request as the host projects it.
Project(Reply<R::Projection>),
Custom(R::Op),
BeforeSend {
wire: Box<WireRequest>,
context: Box<RequestContext>,
reply: Reply<WireRequest>,
},
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<Demand>),
/// The next chunk of an open stream, answered once the caller asks for the one after.
Deliver(R::Chunk),
}
pub enum HostResult<R: Route> {
Route(R::OpResult),
BeforeSend(Box<WireRequest>),
Emitted,
Demand(Demand),
Deliver(R::Chunk, Reply<Demand>),
}
/// Whether the caller of a streamed call still reads it.
@ -39,10 +38,13 @@ pub enum HostStep<V, S> {
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<R: Route>: Send + Sync {
fn route(&self, op: R::Op) -> impl Future<Output = Result<R::OpResult, R::Error>> + Send;
pub trait Host<R: Protocol>: Send + Sync {
fn project(&self) -> impl Future<Output = Result<R::Projection, R::Error>> + Send;
/// Answers `op` through its reply, or fails the call.
fn custom_op(&self, op: R::Op) -> impl Future<Output = Result<(), R::Error>> + Send;
fn before_send(
&self,

7
litellm-rust/crates/host/src/lib.rs
View file

@ -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;

27
litellm-rust/crates/host/src/machine/auth.rs
View file

@ -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<ResolvedCredential>;
/// A protocol whose host can mint credentials on the call's behalf.
pub trait TokenProtocol: Protocol {
fn acquire_token_op(reply: Reply<ResolvedCredential>) -> 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<R: Route> {
pub struct HostTokenProvider<R: Protocol> {
channel: HostChannel<R>,
}
impl<R: Route> std::fmt::Debug for HostTokenProvider<R> {
impl<R: Protocol> std::fmt::Debug for HostTokenProvider<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("HostTokenProvider")
}
@ -24,7 +23,7 @@ impl<R: Route> std::fmt::Debug for HostTokenProvider<R> {
impl<R> HostTokenProvider<R>
where
R: TokenRoute,
R: TokenProtocol,
R::Error: From<MachineFault> + std::fmt::Display,
{
pub fn handle(channel: HostChannel<R>) -> TokenProviderHandle {
@ -34,19 +33,15 @@ where
impl<R> TokenProvider for HostTokenProvider<R>
where
R: TokenRoute,
R: TokenProtocol,
R::Error: From<MachineFault> + 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()))
})
}
}

137
litellm-rust/crates/host/src/machine/call_machine.rs Normal file
View file

@ -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<R> =
Pin<Box<dyn Future<Output = Result<<R as Protocol>::Response, <R as Protocol>::Error>> + Send>>;
/// The provider side of the machine: how the in-flight call reaches its host.
pub struct HostChannel<R: Protocol> {
co: Co<HostOp<R>>,
}
impl<R: Protocol> Clone for HostChannel<R> {
fn clone(&self) -> Self {
Self {
co: self.co.clone(),
}
}
}
impl<R: Protocol> HostChannel<R>
where
R::Error: From<MachineFault>,
{
async fn yield_<A: Send>(
&self,
ask: impl FnOnce(Reply<A>) -> HostOp<R> + Send,
) -> Result<A, R::Error> {
self.co
.yield_(ask)
.await
.map_err(|_| MachineFault::Abandoned.into())
}
pub async fn project(&self) -> Result<R::Projection, R::Error> {
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<A: Send>(
&self,
ask: impl FnOnce(Reply<A>) -> R::Op + Send,
) -> Result<A, R::Error> {
self.yield_(|reply| HostOp::Custom(ask(reply))).await
}
pub async fn before_send(
&self,
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, R::Error> {
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<Demand, R::Error> {
self.yield_(|reply| HostOp::Open(head, reply)).await
}
pub async fn deliver(&self, chunk: R::Chunk) -> Result<Demand, R::Error> {
self.yield_(|reply| HostOp::Deliver(chunk, reply)).await
}
}
type CallCoroutine<R> =
Coroutine<HostOp<R>, Result<<R as Protocol>::Response, <R as Protocol>::Error>>;
pub struct CallMachine<R: Protocol> {
coroutine: CallCoroutine<R>,
}
impl<R: Protocol> CallMachine<R>
where
R::Error: From<MachineFault>,
{
pub fn new(execute: impl FnOnce(HostChannel<R>) -> ExecuteFuture<R> + Send + 'static) -> Self {
Self {
coroutine: Coroutine::new(|co| execute(HostChannel { co })),
}
}
}
impl<R: Protocol> Machine for CallMachine<R>
where
R::Error: From<MachineFault>,
{
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<R::Error>) -> Interrupted<'_, Self> {
self.coroutine.cancel();
Box::pin(async move { Err(failure.into_error()) })
}
}

32
litellm-rust/crates/host/src/machine/mod.rs
View file

@ -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<R: Route, C> {
pub enum MachineStep<R: Protocol, C> {
Host(HostOp<R>),
Complete(C),
}
@ -19,8 +19,8 @@ pub type Step<'a, M> = Pin<
Box<
dyn Future<
Output = Result<
MachineStep<<M as Machine>::Route, <M as Machine>::Complete>,
<<M as Machine>::Route as Route>::Error,
MachineStep<<M as Machine>::Protocol, <M as Machine>::Complete>,
<<M as Machine>::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<<M as Machine>::Complete, <<M as Machine>::Route as Route>::Error>,
Output = Result<
<M as Machine>::Complete,
<<M as Machine>::Protocol as Protocol>::Error,
>,
> + Send
+ 'a,
>,
@ -51,19 +54,18 @@ impl<E> HostFailure<E> {
}
/// 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<HostResult<Self::Route>>) -> 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<<Self::Route as Route>::Error>,
failure: HostFailure<<Self::Protocol as Protocol>::Error>,
) -> Interrupted<'_, Self>;
}

199
litellm-rust/crates/host/src/machine/route_machine.rs
View file

@ -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<R> =
Pin<Box<dyn Future<Output = Result<<R as Route>::Response, <R as Route>::Error>> + Send>>;
struct PendingOp<R: Route> {
op: HostOp<R>,
reply: oneshot::Sender<HostResult<R>>,
}
/// The provider side of the machine: how the in-flight call reaches its host.
pub struct HostChannel<R: Route> {
ops: mpsc::UnboundedSender<PendingOp<R>>,
}
impl<R: Route> Clone for HostChannel<R> {
fn clone(&self) -> Self {
Self {
ops: self.ops.clone(),
}
}
}
impl<R: Route> HostChannel<R>
where
R::Error: From<MachineFault>,
{
async fn invoke(&self, op: HostOp<R>) -> Result<HostResult<R>, 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<R::OpResult, R::Error> {
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<WireRequest, R::Error> {
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<Demand, R::Error> {
self.demand(HostOp::Open(head)).await
}
pub async fn deliver(&self, chunk: R::Chunk) -> Result<Demand, R::Error> {
self.demand(HostOp::Deliver(chunk)).await
}
async fn demand(&self, op: HostOp<R>) -> Result<Demand, R::Error> {
match self.invoke(op).await? {
HostResult::Demand(demand) => Ok(demand),
_ => Err(MachineFault::Mismatch.into()),
}
}
}
enum Execution<R: Route> {
Unstarted(Box<dyn FnOnce(HostChannel<R>) -> ExecuteFuture<R> + Send>),
Running(ExecuteFuture<R>),
Done,
}
pub struct RouteMachine<R: Route> {
execution: Execution<R>,
ops: mpsc::UnboundedReceiver<PendingOp<R>>,
channel: HostChannel<R>,
reply: Option<oneshot::Sender<HostResult<R>>>,
}
impl<R: Route> RouteMachine<R>
where
R::Error: From<MachineFault>,
{
pub fn new(execute: impl FnOnce(HostChannel<R>) -> ExecuteFuture<R> + 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<HostResult<R>>,
) -> Result<MachineStep<R, R::Response>, 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<R: Route> Machine for RouteMachine<R>
where
R::Error: From<MachineFault>,
{
type Route = R;
type Complete = R::Response;
fn resume(&mut self, result: Option<HostResult<R>>) -> Step<'_, Self> {
Box::pin(self.step(result))
}
fn interrupt(&mut self, failure: HostFailure<R::Error>) -> Interrupted<'_, Self> {
self.reply = None;
self.execution = Execution::Done;
Box::pin(async move { Err(failure.into_error()) })
}
}

17
litellm-rust/crates/host/src/protocol.rs Normal file
View file

@ -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;
}

14
litellm-rust/crates/host/src/route.rs
View file

@ -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;
}

147
litellm-rust/crates/host/src/run.rs
View file

@ -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<M, H>(mut machine: M, host: &H) -> Result<M::Complete, <M::Route as Route>::Error>
pub async fn run<M, H>(
mut machine: M,
host: &H,
) -> Result<M::Complete, <M::Protocol as Protocol>::Error>
where
M: Machine,
H: Host<M::Route>,
H: Host<M::Protocol>,
{
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<R: Protocol, H: Host<R>>(host: &H, op: HostOp<R>) -> 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<HostResult<Unit>>) -> 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<MachineFault> for &'static str {
fn from(_: MachineFault) -> Self {
"machine fault"
}
}
@ -100,12 +96,21 @@ mod tests {
}
impl Host<Unit> 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<Unit> {
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<Vec<f64>>);
impl Host<Unit> 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]);

1
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -118,7 +118,6 @@ impl From<litellm_host::machine::MachineFault> 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(),
})
}
}

11
litellm-rust/crates/python-bridge/src/logger/machine.rs
View file

@ -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<M> LoggedMachine<M> {
}
impl<M: Machine> Machine for LoggedMachine<M> {
type Route = M::Route;
type Protocol = M::Protocol;
type Complete = M::Complete;
fn resume(&mut self, result: Option<HostResult<Self::Route>>) -> 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<<Self::Route as Route>::Error>,
failure: HostFailure<<Self::Protocol as Protocol>::Error>,
) -> Interrupted<'_, Self> {
let logger = self.logger.get_or_init(|| Python::attach(super::capture));
Box::pin(logger.instrument(logger.scope(|| self.machine.interrupt(failure))))

13
litellm-rust/crates/python-bridge/src/logger/tests.rs
View file

@ -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<HostResult<Self>>) -> 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<Bound<'_, PyAny>> {
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

33
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -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<PyErr> {
/// 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<PyAny>,
}
impl MessagesRouteHost {
impl MessagesPythonHost {
pub(super) fn new(request: Py<PyAny>) -> Self {
Self { request }
}
fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesCall> {
fn projection(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesCall> {
let request = self.request.bind(py);
let argument = |name: &str| -> PyResult<Option<Bound<'_, PyAny>>> {
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<MessagesOpResult, InvokeError<Error>> {
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<MessagesCall, InvokeError<Error>> {
self.projection(py, arguments)
.map_err(|error| InvokeError::Python(self.map_failure(py, error)))
}
fn invoke(&mut self, _: Python<'_>, op: Infallible) -> Result<(), InvokeError<Error>> {
match op {}
}
fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult<Py<PyAny>> {

4
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -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,
)
}

247
litellm-rust/crates/python-bridge/src/routes/ocr/document.rs
View file

@ -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<PyAny>,
name: Option<String>,
/// 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<String>,
},
}
impl PythonFileReader {
pub(super) fn read(&self, py: Python<'_>) -> PyResult<OcrFileContent> {
let value = self.reader.bind(py).call0()?;
let bytes = if value.is_instance_of::<PyString>() {
Bytes::from(value.extract::<String>()?)
} else if value.is_instance_of::<PyBytes>() {
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<OcrDocumentInput> {
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<Bytes> {
if value.is_exact_instance_of::<PyBytes>() {
return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
}
Ok(Bytes::copy_from_slice(
value.extract::<PyBackedBytes>()?.as_ref(),
))
}
pub(super) struct FileDocumentInput {
pub input: OcrDocumentInput,
pub reader: Option<PythonFileReader>,
}
impl FromPyObject<'_, '_> for FileDocumentInput {
@ -87,51 +67,31 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
}
static PATH_LIKE: PyOnceLock<Py<PyType>> = PyOnceLock::new();
if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? {
return Ok(Self {
input: OcrDocumentInput::Path {
path: file.extract::<PathBuf>()?,
mime_type,
},
reader: None,
});
return Ok(Self::Ready(OcrDocumentInput::Path {
path: file.extract::<PathBuf>()?,
mime_type,
}));
}
if file.is_instance_of::<PyBytes>() {
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::<String>())
.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::<PyValueError>(py));
assert!(error.to_string().contains("bare str"));
let error = py
.eval(c"{'file': object()}", None, None)
.unwrap()
.extract::<FileDocumentInput>()
.err()
.unwrap();
assert!(error.is_instance_of::<PyValueError>(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::<FileDocumentInput>()
.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::<PyTypeError>(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::<FileDocumentInput>()
.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");
}
}

100
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -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<PyAny>,
data: OcrHostData,
}
impl OcrRouteHost {
impl OcrPythonHost {
pub(super) fn new(request: Py<PyAny>) -> Self {
Self {
request,
@ -42,14 +43,6 @@ impl OcrRouteHost {
}
}
fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::types::OcrFileContent> {
self.handles()?
.reader
.as_ref()
.ok_or_else(missing_state)?
.read(py)
}
fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
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<OcrOpResult> {
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<OcrProjection> {
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<OcrOpResult, InvokeError<Error>> {
self.answer(py, arguments, op)
) -> Result<OcrProjection, InvokeError<Error>> {
self.projection(py, arguments)
.map_err(|error| InvokeError::Python(self.map_failure(py, error)))
}
fn invoke(&mut self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError<Error>> {
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<PyAny>> {
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::<PyDict>()
.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 = || {

4
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -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,
)
}

110
litellm-rust/crates/python-bridge/src/routes/ocr/project.rs
View file

@ -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<PythonFileReader>,
pub azure_ad_token_provider: Option<CallerTokenProvider>,
pub provider: &'static str,
}
@ -104,13 +100,11 @@ impl ProjectedDocument {
Ok(Self::File(document.extract()?))
}
fn into_parts(self) -> PyResult<(OcrDocumentInput, Option<PythonFileReader>)> {
/// Reads a file-like document now, so it runs after every other argument was read.
fn resolve(self, py: Python<'_>) -> PyResult<OcrDocumentInput> {
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<PythonFileReader>)> {
ProjectedDocument::project(document)?.into_parts()
fn project_document(document: &Bound<'_, PyAny>) -> PyResult<OcrDocumentInput> {
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::<PyDict>()
.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::<usize>()
.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<String> = document.getattr("reads").unwrap().extract().unwrap();
assert_eq!(reads, ["type", "mime_type", "file"]);