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

refactor(ocr): move file preparation from the python bridge into litellm-core

Browse source 
Delete litellm/ocr/input.py and the native _ocr_file_document, _ocr_upload_document
and _ocr_mime_type helpers. File documents now project to a typed OcrDocumentInput
and the core lifecycle reads local paths, encodes bytes and asks the host to read
file-like objects through a ReadDocument operation before the provider request

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-16 21:02:30 +00:00
parent 79fc5153d3
commit c621435ef7
22 changed files with 810 additions and 517 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -1951,6 +1951,7 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"bytes",
"criterion",
"futures-util",
"litellm-auth",

13
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -54,9 +54,16 @@ impl OcrClient {
match call.resume(result.take()).await? {
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().ok_or_else(|| {
Error::InvalidRequest("OCR request was already projected".into())
})?),
Box::new(
request
.take()
.ok_or_else(|| {
Error::InvalidRequest(
"OCR request was already projected".into(),
)
})?
.into(),
),
false,
))))
}

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

@ -1,3 +1,6 @@
use std::io::Read;
use std::path::Path;
use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::mime::Mime;
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
@ -5,12 +8,52 @@ use reqwest::Url;
use serde_json::Map;
use super::error::{OcrError, OcrRequestError, OcrResponseError};
use super::types::{OcrConnection, OcrDocument};
use super::types::{OcrConnection, OcrDocument, OcrDocumentInput};
use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
use crate::media::Error as MediaError;
use crate::media::{DownloadPolicy, MediaFetcher};
use crate::transport::Error as TransportError;
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::Error> {
match input {
OcrDocumentInput::Document(document) => Ok(document),
OcrDocumentInput::Path { path, mime_type } => {
read_path_document(&path, mime_type.as_deref())
}
OcrDocumentInput::Bytes {
bytes,
file_name,
mime_type,
} => Ok(encode_file_document(
&bytes,
file_name.as_deref(),
mime_type.as_deref(),
)?),
OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest(
"OCR file reader was not read by the host".into(),
)),
}
}
pub fn read_path_document(
path: &Path,
mime_type: Option<&str>,
) -> Result<OcrDocument, super::Error> {
let mut bytes = Vec::new();
std::fs::File::open(path)
.and_then(|file| {
file.take(OCR_INLINE_MAX_BYTES as u64 + 1)
.read_to_end(&mut bytes)
})
.map_err(|source| super::Error::FileRead {
path: path.to_owned(),
kind: source.kind(),
message: source.to_string(),
})?;
let name = path.file_name().map(|name| name.to_string_lossy());
Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?)
}
pub fn encode_file_document(
bytes: &[u8],
file_name: Option<&str>,
@ -75,18 +118,6 @@ pub fn mime_type_for_name(name: &str) -> &'static str {
}
}
pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str {
match content_type
.and_then(|value| value.split(';').next())
.map(str::trim)
{
Some(value) if !value.is_empty() && value != "application/octet-stream" => value,
_ => file_name
.map(mime_type_for_name)
.unwrap_or("application/octet-stream"),
}
}
pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
impl<'a> InlineDocument<'a> {
@ -230,24 +261,65 @@ mod tests {
}
#[test]
fn upload_mime_mapping_matches_python() {
fn path_documents_are_read_and_named_by_core() {
let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("scan.png");
std::fs::write(&path, b"abc").unwrap();
assert_eq!(
upload_mime_type(Some("report.pdf"), Some("application/octet-stream")),
"application/pdf"
);
assert_eq!(upload_mime_type(Some("image.png"), None), "image/png");
assert_eq!(upload_mime_type(None, None), "application/octet-stream");
assert_eq!(
upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")),
"application/pdf"
prepare_document(OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
})
.unwrap(),
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::new(),
}
);
assert_eq!(
upload_mime_type(
Some("img.png"),
Some("image/png; charset=utf-8; boundary=something")
),
"image/png"
prepare_document(OcrDocumentInput::Path {
path: path.clone(),
mime_type: Some("application/pdf".into()),
})
.unwrap(),
document("data:application/pdf;base64,YWJj")
);
std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap();
assert_eq!(
prepare_document(OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
}),
Err(OcrRequestError::InlineDocumentTooLarge.into())
);
std::fs::remove_dir_all(&dir).unwrap();
let missing = dir.join("missing.pdf");
let Err(super::super::Error::FileRead { path, kind, .. }) =
prepare_document(OcrDocumentInput::Path {
path: missing.clone(),
mime_type: None,
})
else {
panic!("missing paths must surface a file read error");
};
assert_eq!(path, missing);
assert_eq!(kind, std::io::ErrorKind::NotFound);
}
#[test]
fn byte_documents_are_encoded_and_host_readers_must_be_read_first() {
assert_eq!(
prepare_document(OcrDocumentInput::Bytes {
bytes: b"abc".as_slice().into(),
file_name: Some("scan.pdf".into()),
mime_type: None,
})
.unwrap(),
document("data:application/pdf;base64,YWJj")
);
assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err());
}
#[test]

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

@ -50,6 +50,12 @@ pub enum Error {
Connect(String),
#[error("routing error: {0}")]
Routing(String),
#[error("Failed to read OCR file {}: {message}", path.display())]
FileRead {
path: std::path::PathBuf,
kind: std::io::ErrorKind,
message: String,
},
/// The request is outside the surface this route covers in Rust. Hosts that
/// keep a reference implementation treat this as "fall back", not "fail".
#[error("unsupported by the rust path: {0}")]

50
litellm-rust/crates/core/src/ocr/lifecycle.rs
View file

@ -9,6 +9,7 @@ use super::hooks::{
OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
OcrPreCallRequest,
};
use super::types::{OcrDocumentInput, OcrFileContent};
use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
use crate::call_lifecycle::host::{
HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
@ -52,6 +53,7 @@ impl OcrAdmission {
#[derive(Clone, Debug)]
pub enum OcrHostOperation {
ProjectRequest,
ReadDocument,
Lifecycle(HostPhase),
ConstructResponse(Arc<LiteLLMOcrResponse>),
MapFailure(Error),
@ -83,7 +85,8 @@ impl OcrHostOperation {
}
pub enum OcrHostResult {
Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
Request(Result<(Box<LiteLLMOcrRequest<OcrDocumentInput>>, bool), Error>),
Document(Result<OcrFileContent, Error>),
Lifecycle(Result<(), HostFailure<Error>>),
AzureAdToken(Result<ResolvedCredential, AuthError>),
PreCall(Result<OcrPreCallRequest, Error>),
@ -313,7 +316,7 @@ struct PendingOperation {
struct OcrExecution {
client: Option<OcrClient>,
request: Option<LiteLLMOcrRequest>,
request: Option<LiteLLMOcrRequest<OcrDocumentInput>>,
operations_tx: mpsc::UnboundedSender<PendingOperation>,
operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
pending_result: Option<oneshot::Sender<OcrHostResult>>,
@ -397,12 +400,14 @@ impl OcrExecution {
},
)));
}
request.hooks = Arc::new(ProtocolHooks {
let hooks = Arc::new(ProtocolHooks {
operations: self.operations_tx.clone(),
intercepts_requests,
terminal: self.terminal.clone(),
});
request.hooks = hooks.clone();
self.execution = Some(tokio::spawn(async move {
let request = prepare_request_document(request, &hooks).await?;
perform_ocr_request(&client, request).await
}));
}
@ -423,6 +428,39 @@ impl OcrExecution {
}
}
async fn prepare_request_document(
request: LiteLLMOcrRequest<OcrDocumentInput>,
hooks: &ProtocolHooks,
) -> Result<LiteLLMOcrRequest, Error> {
let request = match &request.document {
OcrDocumentInput::HostReader { mime_type } => {
let mime_type = mime_type.clone();
let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? {
OcrHostResult::Document(result) => result?,
_ => {
return Err(Error::InvalidRequest(
"invalid OCR document read host result".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);
}
tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document))
.await
.map_err(|error| {
Error::InvalidRequest(format!("OCR document preparation task failed: {error}"))
})?
}
impl Drop for OcrExecution {
fn drop(&mut self) {
if let Some(execution) = &self.execution {
@ -567,6 +605,9 @@ impl OcrHost for NoopOcrHost {
OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
Error::InvalidRequest("OCR host has no request projection".into()),
)),
OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
Error::InvalidRequest("OCR host has no document reader".into()),
)),
OcrHostOperation::Lifecycle(_)
| OcrHostOperation::ConstructResponse(_)
| OcrHostOperation::MapFailure(_)
@ -602,6 +643,9 @@ impl OcrHost for OcrHookHost {
OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
Error::InvalidRequest("OCR hook host has no request projection".into()),
)),
OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
Error::InvalidRequest("OCR hook host has no document reader".into()),
)),
OcrHostOperation::Success {
context,
response,

7
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -13,12 +13,15 @@ pub mod types;
pub mod wire;
pub use client::{OcrClient, ocr};
pub use document::{encode_file_document, mime_type_for_name, upload_mime_type};
pub use document::{encode_file_document, mime_type_for_name, read_path_document};
pub use lifecycle::{
NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
};
pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
pub use types::{
LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrDocumentInput,
OcrFileContent,
};
#[cfg(test)]
#[path = "../../tests/azure_ai_ocr.rs"]

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

@ -1,7 +1,10 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
@ -50,6 +53,35 @@ impl OcrDocument {
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum OcrDocumentInput {
Document(OcrDocument),
Path {
path: PathBuf,
mime_type: Option<String>,
},
Bytes {
bytes: Bytes,
file_name: Option<String>,
mime_type: Option<String>,
},
HostReader {
mime_type: Option<String>,
},
}
impl From<OcrDocument> for OcrDocumentInput {
fn from(document: OcrDocument) -> Self {
Self::Document(document)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OcrFileContent {
pub bytes: Bytes,
pub file_name: Option<String>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OcrResponseFormat {
@ -89,9 +121,9 @@ impl Default for OcrConnection {
}
}
pub struct LiteLLMOcrRequest {
pub struct LiteLLMOcrRequest<D = OcrDocument> {
pub model: String,
pub document: OcrDocument,
pub document: D,
pub connection: OcrConnection,
pub hooks: Arc<dyn OcrHooks>,
pub litellm_call_id: Option<String>,
@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest {
pub(crate) adapter: OcrAdapterKind,
}
impl LiteLLMOcrRequest {
impl<D> LiteLLMOcrRequest<D> {
pub fn new(
model: String,
document: OcrDocument,
document: D,
custom_llm_provider: Option<&str>,
optional_params: Map<String, Value>,
) -> Result<Self, Error> {
@ -151,6 +183,36 @@ impl LiteLLMOcrRequest {
..self
}
}
pub fn map_document<T, E>(
self,
map: impl FnOnce(D) -> Result<T, E>,
) -> Result<LiteLLMOcrRequest<T>, E> {
Ok(LiteLLMOcrRequest {
model: self.model,
document: map(self.document)?,
connection: self.connection,
hooks: self.hooks,
litellm_call_id: self.litellm_call_id,
optional_params: self.optional_params,
input_sources: self.input_sources,
azure_ad_token_provider: self.azure_ad_token_provider,
adapter: self.adapter,
})
}
pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
let Ok(request) = self.map_document(|_| Ok::<T, Infallible>(document));
request
}
}
impl From<LiteLLMOcrRequest> for LiteLLMOcrRequest<OcrDocumentInput> {
fn from(request: LiteLLMOcrRequest) -> Self {
let Ok(request) = request
.map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document)));
request
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

43
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -68,9 +68,9 @@ pub struct DecodedOcrResponse<T> {
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OcrWireRequest {
pub struct OcrWireRequest<D = Value> {
pub model: String,
pub document: Value,
pub document: D,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
@ -141,10 +141,34 @@ pub fn consumed_optional_params(
}
pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
let OcrWireRequest {
model,
document,
api_key,
api_base,
custom_llm_provider,
extra_headers,
optional_params,
input_sources,
timeout_seconds,
} = wire;
decode_request_input(OcrWireRequest {
model,
document: decode_document(document)?,
api_key,
api_base,
custom_llm_provider,
extra_headers,
optional_params,
input_sources,
timeout_seconds,
})
}
pub fn decode_request_input<D>(wire: OcrWireRequest<D>) -> Result<LiteLLMOcrRequest<D>, Error> {
let api_key_source = source_for(&wire.input_sources, "api_key");
let api_base_source = source_for(&wire.input_sources, "api_base");
let extra_headers_source = source_for(&wire.input_sources, "extra_headers");
let document = decode_document(wire.document)?;
let headers = wire
.extra_headers
.unwrap_or_default()
@ -183,7 +207,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
.unwrap_or(defaults.max_response_bytes);
let request = LiteLLMOcrRequest::new(
wire.model,
document,
wire.document,
wire.custom_llm_provider.as_deref(),
wire.optional_params
.into_iter()
@ -209,14 +233,14 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
})
}
fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
let kind = value.get("type").and_then(Value::as_str);
let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none()
|| matches!(kind, Some("image_url")) && value.get("image_url").is_none();
if missing_url {
return Err(OcrRequestError::MissingDocumentUrl);
return Err(OcrRequestError::MissingDocumentUrl.into());
}
decode_request_value(value, "document")
Ok(decode_request_value(value, "document")?)
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
@ -334,10 +358,7 @@ mod tests {
serde_json::json!({"type": "document_url"}),
serde_json::json!({"type": "image_url"}),
] {
assert_eq!(
decode_document(document),
Err(OcrRequestError::MissingDocumentUrl)
);
assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl));
}
}
}

154
litellm-rust/crates/core/tests/ocr.rs
View file

@ -348,13 +348,14 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() {
}
OcrHostOperation::ProjectRequest => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().unwrap()),
Box::new(request.take().unwrap().into()),
false,
))))
}
OcrHostOperation::AcquireAzureAdToken => {
panic!("test request has no token provider")
}
OcrHostOperation::ReadDocument => panic!("test request has no file reader"),
OcrHostOperation::PreCall(request) => {
phases.push("pre");
result = Some(OcrHostResult::PreCall(if failure_phase == "pre" {
@ -405,7 +406,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure()
match call.resume(result.take()).await {
Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().unwrap()),
Box::new(request.take().unwrap().into()),
false,
))));
}
@ -467,9 +468,10 @@ async fn direct_native_host_drives_the_same_state_machine() {
_ => panic!("unexpected OCR operation"),
});
result = Some(match operation {
OcrHostOperation::ProjectRequest => {
OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
}
OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
Box::new(request.take().unwrap().into()),
false,
))),
operation => host.invoke(operation).await,
});
}
@ -501,6 +503,137 @@ async fn direct_native_host_drives_the_same_state_machine() {
));
}
async fn drive_native_file_call(
request: super::LiteLLMOcrRequest<super::OcrDocumentInput>,
content: Result<super::OcrFileContent, crate::ocr::Error>,
) -> (Result<super::LiteLLMOcrResponse, crate::ocr::Error>, usize) {
let NativeOutcome::Completed(mut call) =
OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
else {
panic!("supported call declined")
};
let mut request = Some(request);
let mut content = Some(content);
let mut result = None;
let mut reads = 0;
let outcome = loop {
match call.resume(result.take()).await {
Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().unwrap()),
false,
))));
}
Ok(OcrCallStep::Host(OcrHostOperation::ReadDocument)) => {
reads += 1;
result = Some(OcrHostResult::Document(content.take().unwrap()));
}
Ok(OcrCallStep::Host(operation)) => result = Some(NoopOcrHost.invoke(operation).await),
Ok(OcrCallStep::Complete(response)) => break Ok(response),
Err(error) => break Err(error),
}
};
(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(
super::OcrDocumentInput::HostReader {
mime_type: Some("application/pdf".into()),
},
);
let (response, reads) = drive_native_file_call(
request,
Ok(super::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() {
let (base, seen, _server) = mock_server(vec![]).await;
let request = wire_request("mistral/model", &base, json!({}));
let failure = crate::ocr::Error::InvalidRequest("reader exploded".into());
let (response, reads) = drive_native_file_call(
request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
Err(failure.clone()),
)
.await;
assert_eq!(response.unwrap_err(), failure);
assert_eq!(reads, 1);
let request = wire_request("mistral/model", &base, json!({}));
let (response, _) = drive_native_file_call(
request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
Ok(super::OcrFileContent {
bytes: Default::default(),
file_name: None,
}),
)
.await;
assert!(matches!(
response.unwrap_err(),
crate::ocr::Error::InvalidRequest(_)
));
assert!(seen.lock().unwrap().is_empty());
}
#[tokio::test]
async fn path_documents_are_read_by_core_without_a_host_operation() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"path"}]
}))])
.await;
let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("scan.png");
std::fs::write(&path, b"abc").unwrap();
let request = wire_request("mistral/model", &base, json!({})).with_document(
super::OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
},
);
let (response, reads) = drive_native_file_call(
request,
Err(crate::ocr::Error::InvalidRequest("unused".into())),
)
.await;
server.await.unwrap();
std::fs::remove_dir_all(&dir).unwrap();
assert_eq!(response.unwrap().pages[0]["markdown"], "path");
assert_eq!(reads, 0);
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(super::OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
}),
Err(crate::ocr::Error::InvalidRequest("unused".into())),
)
.await;
assert!(matches!(
response.unwrap_err(),
crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path
));
assert!(seen.lock().unwrap().is_empty());
}
#[tokio::test]
async fn public_finalization_failure_never_dispatches_success_or_replays_provider() {
use crate::call_lifecycle::host::{HostFailure, HostPhase};
@ -543,9 +676,10 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide
| OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => {
panic!("finalization failure used provider/success dispatch")
}
OcrHostOperation::ProjectRequest => {
OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
}
OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
Box::new(request.take().unwrap().into()),
false,
))),
operation => host.invoke(operation).await,
});
}
@ -582,7 +716,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption
OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break,
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().unwrap()),
Box::new(request.take().unwrap().into()),
false,
))))
}
@ -799,7 +933,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_
_ = entered.notified() => break,
step = call.resume(result.take()) => {
result = Some(match step.unwrap() {
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))),
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap().into()), false))),
OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await,
OcrCallStep::Complete(_) => panic!("pending provider completed"),
});

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

@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"]
panic-test = []
[dependencies]
bytes.workspace = true
futures-util.workspace = true
litellm-core.workspace = true
litellm-auth.workspace = true

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

@ -1,97 +1,56 @@
use std::io::Read;
use std::path::PathBuf;
use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError};
use bytes::Bytes;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedBytes;
#[cfg(test)]
use pyo3::types::PyDict;
use pyo3::types::{PyBytes, PyString};
use litellm_core::constants::OCR_INLINE_MAX_BYTES;
use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type};
use litellm_python_interop::to_py_preserving_errors;
use litellm_core::ocr::{OcrDocumentInput, OcrFileContent};
enum FileBytes {
Python(PyBackedBytes),
Native(Vec<u8>),
#[derive(Debug)]
pub(super) struct PythonFileReader {
reader: Py<PyAny>,
name: Option<String>,
}
impl AsRef<[u8]> for FileBytes {
fn as_ref(&self) -> &[u8] {
match self {
Self::Python(bytes) => bytes,
Self::Native(bytes) => bytes,
}
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(),
})
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reader)
}
}
fn read_file_input(
py: Python<'_>,
file: &Bound<'_, PyAny>,
) -> PyResult<(FileBytes, Option<String>)> {
if file.is_instance_of::<PyString>() {
return Err(PyValueError::new_err(
"OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
));
fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult<Bytes> {
if value.is_exact_instance_of::<PyBytes>() {
return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
}
if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
let path: PathBuf = file.extract()?;
let name = path
.file_name()
.map(|value| value.to_string_lossy().into_owned());
let bytes = py
.detach(|| {
let mut bytes = Vec::new();
std::fs::File::open(&path)?
.take(OCR_INLINE_MAX_BYTES as u64 + 1)
.read_to_end(&mut bytes)?;
Ok::<_, std::io::Error>(bytes)
})
.map_err(|error| {
if error.kind() == std::io::ErrorKind::NotFound {
PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
} else {
error.into()
}
})?;
return Ok((FileBytes::Native(bytes), name));
}
if file.is_instance_of::<PyBytes>() {
return Ok((FileBytes::Python(file.extract()?), None));
}
let reader = file
.getattr_opt("read")?
.filter(|value| value.is_callable());
let Some(reader) = reader else {
return 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()?;
let value = reader.call0()?;
let bytes = if value.is_instance_of::<PyString>() {
FileBytes::Native(value.extract::<String>()?.into_bytes())
} else if value.is_instance_of::<PyBytes>() {
FileBytes::Python(value.extract()?)
} else {
return Err(PyTypeError::new_err(format!(
"OCR file read must return bytes or str, got {}",
value.get_type(),
)));
};
Ok((bytes, name))
Ok(Bytes::copy_from_slice(
value.extract::<PyBackedBytes>()?.as_ref(),
))
}
pub(super) struct FileDocumentInput {
bytes: FileBytes,
name: Option<String>,
mime_type: Option<String>,
pub input: OcrDocumentInput,
pub reader: Option<PythonFileReader>,
}
impl FromPyObject<'_, '_> for FileDocumentInput {
@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => None,
Err(error) => return Err(error),
};
let missing = || {
PyValueError::new_err(
"document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
)
};
let file = document.get_item("file").map_err(|error| {
if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) {
PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes")
missing()
} else {
error
}
})?;
if file.is_none() {
return Err(missing());
}
if file.is_instance_of::<PyString>() {
return Err(PyValueError::new_err(
"document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
"OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
));
}
let (bytes, name) = read_file_input(py, &file)?;
if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
return Ok(Self {
input: OcrDocumentInput::Path {
path: file.extract::<PathBuf>()?,
mime_type,
},
reader: None,
});
}
if file.is_instance_of::<PyBytes>() {
return Ok(Self {
input: OcrDocumentInput::Bytes {
bytes: extract_bytes(&file)?,
file_name: None,
mime_type,
},
reader: None,
});
}
let reader = file
.getattr_opt("read")?
.filter(|value| value.is_callable());
let Some(reader) = reader else {
return 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 {
bytes,
name,
mime_type,
input: OcrDocumentInput::HostReader { mime_type },
reader: Some(PythonFileReader {
reader: reader.unbind(),
name,
}),
})
}
}
pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult<OcrDocument> {
py.detach(|| {
encode_file_document(
document.bytes.as_ref(),
document.name.as_deref(),
document.mime_type.as_deref(),
)
})
.map_err(|error| PyValueError::new_err(error.to_string()))
}
#[pyfunction]
fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
to_py_preserving_errors(py, &file_document(py, document.extract()?)?)
}
#[pyfunction]
fn _ocr_mime_type(file_name: &str) -> String {
mime_type_for_name(file_name).into()
}
#[pyfunction]
#[pyo3(signature = (file_content, file_name=None, content_type=None))]
fn _ocr_upload_document(
py: Python<'_>,
file_content: &Bound<'_, PyBytes>,
file_name: Option<&str>,
content_type: Option<&str>,
) -> PyResult<Py<PyAny>> {
let bytes: PyBackedBytes = file_content.extract()?;
let document = py
.detach(|| {
encode_file_document(
&bytes,
None,
Some(upload_mime_type(file_name, content_type)),
)
})
.map_err(|error| PyValueError::new_err(error.to_string()))?;
to_py_preserving_errors(py, &document)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?;
module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?;
module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?;
module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?)
}
#[cfg(test)]
mod tests {
use super::*;
use pyo3::types::PyDict;
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
}
#[test]
fn extraction_validates_required_file_and_optional_mime_type() {
@ -196,69 +155,148 @@ mod tests {
let error = document.extract::<FileDocumentInput>().err().unwrap();
assert!(error.is_instance_of::<PyTypeError>(py));
}
let document = py.eval(c"{'file': b'abc'}", None, None).unwrap();
let error = py
.eval(c"{'file': 'scan.pdf'}", None, None)
.unwrap()
.extract::<FileDocumentInput>()
.err()
.unwrap();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("bare str"));
let document = py
.eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None)
.unwrap();
let input: FileDocumentInput = document.extract().unwrap();
assert_eq!(input.bytes.as_ref(), b"abc");
assert_eq!(input.name, None);
assert_eq!(input.mime_type, None);
assert!(input.reader.is_none());
assert_eq!(
input.input,
OcrDocumentInput::Bytes {
bytes: b"abc".as_slice().into(),
file_name: None,
mime_type: Some("image/png".into()),
}
);
});
}
#[test]
fn extraction_validates_mime_type_before_consuming_file() {
fn paths_and_readers_are_projected_without_io() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"class Reader:
let locals = eval(
py,
c"from pathlib import Path
class Reader:
name = 'scan.png'
def __init__(self):
self.reads = 0
def read(self):
self.reads += 1
return b'abc'
reader = Reader()
document = {'file': reader, 'mime_type': 7}",
Some(&locals),
Some(&locals),
)
.unwrap();
document = {'file': reader, 'mime_type': 7}
reader_document = {'file': reader}
path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}",
);
let document = locals.get_item("document").unwrap().unwrap();
let error = document.extract::<FileDocumentInput>().err().unwrap();
assert!(error.is_instance_of::<PyTypeError>(py));
let reads: usize = locals
.get_item("reader")
.unwrap()
.unwrap()
.getattr("reads")
.unwrap()
.extract()
.unwrap();
assert_eq!(reads, 0);
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")
.unwrap()
.unwrap()
.getattr("reads")
.unwrap()
.extract::<usize>()
.unwrap()
};
assert_eq!(reads(), 0);
let content = input.reader.unwrap().read(py).unwrap();
assert_eq!(reads(), 1);
assert_eq!(
content,
OcrFileContent {
bytes: b"abc".as_slice().into(),
file_name: Some("scan.png".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,
OcrDocumentInput::Path {
path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"),
mime_type: Some("image/png".into()),
}
);
});
}
#[test]
fn extraction_preserves_reader_key_error_identity() {
fn reader_results_are_normalized_and_exceptions_keep_their_identity() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
let locals = eval(
py,
c"failure = KeyError('reader failed')
class Reader:
class Raising:
def read(self):
raise failure
document = {'file': Reader()}",
Some(&locals),
Some(&locals),
)
.unwrap();
let document = locals.get_item("document").unwrap().unwrap();
let error = document.extract::<FileDocumentInput>().err().unwrap();
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"));
});
}
#[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");
}
}

7
litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs
View file

@ -1,4 +1,5 @@
use litellm_core::ocr::Error;
use pyo3::exceptions::{PyFileNotFoundError, PyOSError};
use pyo3::prelude::*;
use crate::errors::{RustUpstreamError, core_error_to_pyerr};
@ -7,6 +8,12 @@ pub(super) fn to_pyerr(error: Error) -> PyErr {
let status = error.http_status_code();
let mapped = match error {
Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
Error::FileRead {
path,
kind: std::io::ErrorKind::NotFound,
..
} => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())),
Error::FileRead { message, .. } => PyOSError::new_err(message),
other => core_error_to_pyerr(other.into()),
};
attach_status(mapped, status)

22
litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs
View file

@ -66,13 +66,27 @@ impl PythonOcrHost {
retained_fields.set_item(name, value)?;
}
}
retained_fields.set_item("document", &self.projected()?.fields.document)?;
let projected = self.projected_mut()?;
let document = match &projected.fields.document {
Some(document) => document.clone_ref(py),
None => to_py(py, &request.document)?,
};
retained_fields.set_item("document", &document)?;
projected.fields.document = Some(document);
projected.retained_fields = Some(retained_fields.unbind());
projected.pre_call = Some((&request).into());
Ok(request)
}
fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::OcrFileContent> {
self.projected()?
.fields
.reader
.as_ref()
.ok_or_else(missing_state)?
.read(py)
}
fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
let provider = self
.projected()?
@ -193,7 +207,7 @@ impl PythonRoute for PythonOcrHost {
let OcrHostData::Unprojected { request } = &self.data else {
return Err(missing_state());
};
let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?;
let projected = project_request(request.bind(py), self.state.kwargs.bind(py))?;
let has_token_provider = projected.fields.azure_ad_token_provider.is_some();
let request = projected.request;
self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost {
@ -205,6 +219,7 @@ impl PythonRoute for PythonOcrHost {
}));
OcrHostResult::Request(Ok((Box::new(request), has_token_provider)))
}
OcrHostOperation::ReadDocument => OcrHostResult::Document(Ok(self.read_document(py)?)),
OcrHostOperation::AcquireAzureAdToken => {
OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?))
}
@ -258,6 +273,9 @@ impl PythonRoute for PythonOcrHost {
OcrHostData::Projected(projected) => {
visit.call(&projected.fields.boundary_request)?;
visit.call(&projected.fields.document)?;
if let Some(reader) = &projected.fields.reader {
reader.traverse(visit)?;
}
visit.call(&projected.fields.api_key)?;
if let Some(provider) = &projected.fields.azure_ad_token_provider {
provider.traverse(visit)?;

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

@ -9,6 +9,5 @@ use pyo3::prelude::*;
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
value::register(module)?;
document::register(module)?;
lifecycle::register(module)
}

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

@ -1,14 +1,15 @@
use std::sync::Arc;
use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request};
use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall};
use litellm_python_interop::{
from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
use litellm_core::ocr::wire::{
OcrWireRequest, consumed_optional_params, decode_document, decode_request_input,
};
use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall, OcrDocumentInput};
use litellm_python_interop::from_py_preserving_errors as from_py;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::{Map, Value};
use super::document::{FileDocumentInput, PythonFileReader};
use super::errors::to_pyerr as ocr_error_to_pyerr;
use super::lifecycle::BridgeOcrHooks;
use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider};
@ -17,7 +18,8 @@ use crate::marshal::{project_optional_fields, python_timeout_seconds, request_in
pub(super) struct ProjectedOcrFields {
pub boundary_request: Py<PyAny>,
pub document: Py<PyAny>,
pub document: Option<Py<PyAny>>,
pub reader: Option<PythonFileReader>,
pub api_key: Py<PyAny>,
pub azure_ad_token_provider: Option<PythonTokenProvider>,
pub provider: &'static str,
@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields {
}
pub(super) struct ProjectedOcrCall {
pub request: LiteLLMOcrRequest,
pub request: LiteLLMOcrRequest<OcrDocumentInput>,
pub fields: ProjectedOcrFields,
}
@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> {
}
enum ProjectedDocument {
File { wire: Value, retained: Py<PyAny> },
File(FileDocumentInput),
Other { wire: Value, retained: Py<PyAny> },
}
impl ProjectedDocument {
fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult<Self> {
fn project(document: &Bound<'_, PyAny>) -> PyResult<Self> {
let kind: String = document.get_item("type")?.extract()?;
if kind != "file" {
return Ok(Self::Other {
@ -93,25 +95,28 @@ impl ProjectedDocument {
retained: document.clone().unbind(),
});
}
let input = document.extract()?;
let encoded = super::document::file_document(py, input)?;
let wire = serde_json::to_value(encoded)
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?;
Ok(Self::File {
retained: to_py(py, &wire)?,
wire,
})
Ok(Self::File(document.extract()?))
}
fn into_parts(self) -> (Value, Py<PyAny>) {
fn into_parts(
self,
) -> PyResult<(
OcrDocumentInput,
Option<Py<PyAny>>,
Option<PythonFileReader>,
)> {
match self {
Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained),
Self::File(FileDocumentInput { input, reader }) => Ok((input, None, reader)),
Self::Other { wire, retained } => Ok((
decode_document(wire).map_err(ocr_error_to_pyerr)?.into(),
Some(retained),
None,
)),
}
}
}
pub(super) fn project_request(
py: Python<'_>,
request: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<ProjectedOcrCall> {
@ -119,8 +124,7 @@ pub(super) fn project_request(
let arguments = OcrArguments { request, kwargs };
let model = arguments.model()?;
let custom_llm_provider = arguments.custom_llm_provider()?;
let (wire_document, retained_document) =
ProjectedDocument::project(py, &arguments.document()?)?.into_parts();
let document = ProjectedDocument::project(&arguments.document()?)?;
let api_key = arguments.api_key()?;
let specs = consumed_optional_params(&model, custom_llm_provider.as_deref())
.map_err(ocr_error_to_pyerr)?;
@ -136,9 +140,10 @@ pub(super) fn project_request(
let azure_ad_token_provider = kwargs
.get_item("azure_ad_token_provider")?
.and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER));
let (document, retained_document, reader) = document.into_parts()?;
let wire = OcrWireRequest {
model,
document: wire_document,
document,
api_key: api_key.extract()?,
api_base: arguments.api_base()?,
custom_llm_provider,
@ -147,13 +152,14 @@ pub(super) fn project_request(
input_sources,
timeout_seconds: arguments.timeout_seconds()?,
};
let request = decode_request(wire).map_err(ocr_error_to_pyerr)?;
let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?;
let provider = request.provider_name();
Ok(ProjectedOcrCall {
request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None),
fields: ProjectedOcrFields {
boundary_request,
document: retained_document,
reader,
api_key: api_key.unbind(),
azure_ad_token_provider,
provider,
@ -197,10 +203,21 @@ mod tests {
}
fn project_document(
py: Python<'_>,
document: &Bound<'_, PyAny>,
) -> PyResult<(Value, Py<PyAny>)> {
ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts)
) -> PyResult<(
OcrDocumentInput,
Option<Py<PyAny>>,
Option<PythonFileReader>,
)> {
ProjectedDocument::project(document)?.into_parts()
}
fn url_document(url: &str) -> OcrDocumentInput {
litellm_core::ocr::OcrDocument::DocumentUrl {
document_url: url.into(),
extra_fields: Map::new(),
}
.into()
}
fn stub_timeout_conversion(py: Python<'_>) {
@ -374,7 +391,7 @@ kwargs = {}
}
#[test]
fn document_reader_mutations_are_visible_to_later_field_reads() {
fn document_readers_are_not_consumed_during_projection() {
Python::initialize();
Python::attach(|py| {
stub_timeout_conversion(py);
@ -406,7 +423,12 @@ kwargs = {}
.unwrap();
let arguments = arguments(&request, &kwargs);
let document = arguments.document().unwrap();
project_document(py, &document).unwrap();
let (input, retained, reader) = project_document(&document).unwrap();
assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None });
assert!(retained.is_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));
});
@ -444,7 +466,7 @@ kwargs = {'api_key': key}
}
#[test]
fn file_documents_are_encoded_and_other_documents_keep_the_python_object() {
fn file_documents_become_typed_inputs_and_other_documents_keep_the_python_object() {
Python::initialize();
Python::attach(|py| {
let file = py
@ -454,13 +476,17 @@ kwargs = {'api_key': key}
None,
)
.unwrap();
let (input, retained, reader) = project_document(&file).unwrap();
assert_eq!(
project_document(py, &file).unwrap().0,
serde_json::json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
})
input,
OcrDocumentInput::Bytes {
bytes: b"%PDF-1.4".as_slice().into(),
file_name: None,
mime_type: Some("application/pdf".into()),
}
);
assert!(retained.is_none());
assert!(reader.is_none());
let original = py
.eval(
@ -469,44 +495,21 @@ kwargs = {'api_key': key}
None,
)
.unwrap();
let (wire, retained) = project_document(py, &original).unwrap();
assert_eq!(
wire,
serde_json::json!({
"type": "document_url",
"document_url": "https://example.com/a.pdf",
})
);
assert!(retained.bind(py).is(&original));
let (input, retained, _) = project_document(&original).unwrap();
assert_eq!(input, url_document("https://example.com/a.pdf"));
assert!(retained.unwrap().bind(py).is(&original));
});
}
#[test]
fn unknown_document_types_reach_existing_downstream_validation() {
fn unknown_document_types_reach_existing_core_validation() {
Python::initialize();
Python::attach(|py| {
let document = py
.eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None)
.unwrap();
let wire_document = project_document(py, &document).unwrap().0;
assert_eq!(
wire_document,
serde_json::json!({"type": "mystery", "mystery": "x"})
);
let error = match decode_request(OcrWireRequest {
model: "mistral/mistral-ocr-latest".into(),
document: wire_document,
api_key: None,
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
input_sources: Default::default(),
timeout_seconds: None,
}) {
Ok(_) => panic!("unknown discriminators belong to core validation"),
Err(error) => error,
};
let error = project_document(&document).unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("document"));
});
}
@ -517,14 +520,14 @@ kwargs = {'api_key': key}
Python::attach(|py| {
let missing = py.eval(c"{}", None, None).unwrap();
assert!(
project_document(py, &missing)
project_document(&missing)
.unwrap_err()
.is_instance_of::<PyKeyError>(py)
);
let non_string = py.eval(c"{'type': 1}", None, None).unwrap();
assert!(
project_document(py, &non_string)
project_document(&non_string)
.unwrap_err()
.is_instance_of::<PyTypeError>(py)
);
@ -540,7 +543,7 @@ document = Document()
",
);
let error =
project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err();
project_document(&locals.get_item("document").unwrap().unwrap()).unwrap_err();
assert!(
error
.value(py)
@ -569,9 +572,9 @@ document = Document()
",
);
let document = locals.get_item("document").unwrap().unwrap();
let (wire, retained) = project_document(py, &document).unwrap();
assert_eq!(wire["type"], "document_url");
assert!(!retained.bind(py).is(&document));
let (input, retained, _) = project_document(&document).unwrap();
assert!(matches!(input, OcrDocumentInput::Bytes { .. }));
assert!(retained.is_none());
let reads: Vec<String> = document.getattr("reads").unwrap().extract().unwrap();
assert_eq!(reads, ["type", "mime_type", "file"]);
});

112
litellm/ocr/input.py
View file

@ -1,112 +0,0 @@
from collections.abc import Mapping
from os import PathLike
from typing import Final, Literal, Protocol, cast # noqa: TID251 # native callables are validated when loaded
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.configuration import rust_ocr_enabled
class FileReader(Protocol):
def read(self) -> bytes | str: ...
class FileDocument(TypedDict):
type: ReadOnly[Literal["file"]]
file: ReadOnly[bytes | PathLike[str] | FileReader]
mime_type: ReadOnly[NotRequired[str]]
class NativeFileDocument(Protocol):
def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ...
class NativeUploadDocument(Protocol):
def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ...
class NativeMimeType(Protocol):
def __call__(self, file_name: str) -> str: ...
_FILE_DOCUMENT: Final = NativeBinding(
"_ocr_file_document",
validate=lambda value: (
cast( # cast-ok: native export owns the callable signature
NativeFileDocument, value
)
if callable(value)
else None
),
)
_UPLOAD_DOCUMENT: Final = NativeBinding(
"_ocr_upload_document",
validate=lambda value: (
cast( # cast-ok: native export owns the callable signature
NativeUploadDocument, value
)
if callable(value)
else None
),
)
_MAX_FILE_BYTES: Final = NativeBinding(
"_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None
)
_MIME_TYPE: Final = NativeBinding(
"_ocr_mime_type",
validate=lambda value: (
cast( # cast-ok: native export owns the callable signature
NativeMimeType, value
)
if callable(value)
else None
),
)
_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
def get_mime_type(file_path: str) -> str:
native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None
if native is None:
from litellm.ocr import legacy
return legacy.get_mime_type(file_path)
return native(file_path)
def get_max_file_bytes() -> int:
limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None
if limit is None:
return _PYTHON_MAX_FILE_BYTES
return limit
def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]:
native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None
if native is None:
from litellm.ocr import legacy
return legacy.convert_file_document_to_url_document(document)
return native(document)
def convert_upload_to_url_document(
file_content: bytes, filename: str | None, content_type: str | None
) -> dict[str, str]:
native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None
if native is None:
from litellm.ocr import legacy
if len(file_content) > _PYTHON_MAX_FILE_BYTES:
raise ValueError("OCR file exceeds the size limit")
content_mime: Final = content_type.split(";")[0].strip() if content_type else None
mime_type: Final = (
legacy.get_mime_type(filename)
if filename and (not content_mime or content_mime == "application/octet-stream")
else content_mime or "application/octet-stream"
)
return legacy.convert_file_document_to_url_document(
{"type": "file", "file": file_content, "mime_type": mime_type}
)
return native(file_content, filename, content_type)

7
litellm/ocr/legacy.py
View file

@ -11,7 +11,7 @@ from collections.abc import Coroutine, Mapping
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
from typing import Final, Protocol, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
import httpx
@ -26,7 +26,6 @@ from litellm.llms.base_llm.ocr.transformation import (
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.ocr.input import FileReader
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -34,6 +33,10 @@ from litellm.utils import ProviderConfigManager, client
base_llm_http_handler: Final = BaseLLMHTTPHandler()
class FileReader(Protocol):
def read(self) -> bytes | str: ...
@dataclass(frozen=True, slots=True)
class _PreparedOCRRequest:
model: str

2
litellm/ocr/main.py
View file

@ -5,7 +5,7 @@ import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import legacy
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.configuration import rust_ocr_enabled
from litellm.rust_bridge.ocr import LiteLLMOcrRequest

17
litellm/proxy/ocr_endpoints/endpoints.py
View file

@ -15,12 +15,13 @@ from litellm.llms.base_llm.ocr.transformation import (
OCRResponse,
parse_ocr_request_format,
)
from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
router: Final = APIRouter()
_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
def _build_document_from_upload(
@ -28,7 +29,15 @@ def _build_document_from_upload(
filename: str | None,
content_type: str | None,
) -> dict[str, str]:
return convert_upload_to_url_document(file_content, filename, content_type)
supplied_mime: Final = content_type.split(";")[0].strip() if content_type else None
mime_type: Final = (
get_mime_type(filename)
if filename and (not supplied_mime or supplied_mime == "application/octet-stream")
else supplied_mime
)
return convert_file_document_to_url_document(
{"type": "file", "file": file_content, "mime_type": mime_type or "application/octet-stream"}
)
def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]:
@ -103,9 +112,11 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]:
# Seek to start in case the file was already partially read by middleware
await uploaded_file.seek(0)
file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1)
file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1)
if not file_content:
raise ValueError("Uploaded file is empty")
if len(file_content) > _MAX_FILE_BYTES:
raise ValueError("OCR file exceeds the size limit")
document: Final = _build_document_from_upload(
file_content=file_content,

13
litellm/rust_bridge/_native.pyi
View file

@ -33,15 +33,6 @@ def aocr(
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
_OCR_MAX_FILE_BYTES: int
def _ocr_upload_document(
file_content: bytes,
file_name: str | None = None,
content_type: str | None = None,
) -> dict[str, str]: ...
def _ocr_file_document(document: Mapping[str, object]) -> dict[str, str]: ...
def _ocr_mime_type(file_name: str) -> str: ...
def _ocr_lifecycle(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
@ -139,15 +130,11 @@ class TokenCounter:
def gil_stats() -> dict[str, int]: ...
__all__ = [
"_OCR_MAX_FILE_BYTES",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
"RustUpstreamError",
"TokenCounter",
"_ocr_file_document",
"_ocr_lifecycle",
"_ocr_mime_type",
"_ocr_upload_document",
"achat_completions",
"amessages",
"aocr",

25
tests/test_litellm/ocr/test_ocr_file_input.py
View file

@ -12,32 +12,16 @@ Tests that:
import base64
import os
import tempfile
from collections.abc import Generator
from io import BytesIO
from pathlib import Path
from typing import Final
from unittest.mock import AsyncMock, MagicMock, Mock
from unittest.mock import AsyncMock, MagicMock
import orjson
import pytest
from starlette.datastructures import FormData
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"])
def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
from litellm.rust_bridge import bindings, configuration
configuration.reset_rust_configuration()
monkeypatch.delenv("LITELLM_RUST", raising=False)
if request.param == "disabled":
monkeypatch.setenv("LITELLM_RUST", "0")
monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled")))
elif request.param == "unavailable":
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
yield
configuration.reset_rust_configuration()
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
class TestGetMimeType:
@ -503,10 +487,9 @@ class TestProxySecurityGuard:
async def test_proxy_upload_stops_reading_at_size_limit() -> None:
from starlette.datastructures import UploadFile
from litellm.ocr.input import get_max_file_bytes
from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form
from litellm.proxy.ocr_endpoints.endpoints import _MAX_FILE_BYTES, _parse_multipart_form
limit: Final = get_max_file_bytes()
limit: Final = _MAX_FILE_BYTES
with tempfile.TemporaryFile() as stream:
stream.truncate(limit * 2)
upload: Final = UploadFile(file=stream, filename="large.pdf")

139
tests/test_litellm_rust/ocr/test_requests.py
View file

@ -518,32 +518,34 @@ async def test_native_ocr_inherits_named_credentials_without_overwriting_argumen
assert ocr_server.requests[0].body["pages"] == [0, 2]
@pytest.mark.parametrize("source", ["sdk", "proxy"])
@pytest.mark.parametrize(
"filename,mime", [("scan.PNG", "image/png"), ("document.pdf", "application/pdf"), ("note.txt", "text/plain")]
"filename,field,mime",
[("scan.PNG", "image_url", "image/png"), ("document.pdf", "document_url", "application/pdf")],
)
def test_ocr_file_helpers_use_native_document_preparation(source: str, filename: str, mime: str) -> None:
def test_native_ocr_infers_mime_type_from_reader_name(
ocr_server: RecordingServer, filename: str, field: str, mime: str
) -> None:
from io import BytesIO
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload
file: Final = BytesIO(b"abc")
file.name = filename
document: Final = (
convert_file_document_to_url_document({"type": "file", "file": file})
if source == "sdk"
else _build_document_from_upload(b"abc", filename, "application/octet-stream; charset=utf-8")
)
field: Final = "image_url" if mime.startswith("image/") else "document_url"
assert get_mime_type(filename) == mime
assert document == {"type": field, field: f"data:{mime};base64,YWJj"}
call_native_ocr(ocr_server, document={"type": "file", "file": file})
assert ocr_server.requests[0].body["document"] == {"type": field, field: f"data:{mime};base64,YWJj"}
def test_native_ocr_encodes_str_reader_results_as_utf8(ocr_server: RecordingServer) -> None:
from io import StringIO
call_native_ocr(ocr_server, document={"type": "file", "file": StringIO("abc"), "mime_type": "text/plain"})
assert ocr_server.requests[0].body["document"] == {
"type": "document_url",
"document_url": "data:text/plain;base64,YWJj",
}
@pytest.mark.parametrize("attribute", ["read", "name"])
def test_native_file_preparation_preserves_property_errors(attribute: str) -> None:
from litellm.ocr.input import convert_file_document_to_url_document
def test_native_file_preparation_preserves_property_errors(ocr_server: RecordingServer, attribute: str) -> None:
ocr_server.expected_requests = 0
failure: Final = LookupError("file property failed")
class File:
@ -555,16 +557,47 @@ def test_native_file_preparation_preserves_property_errors(attribute: str) -> No
def read(self):
return b"abc"
with pytest.raises(LookupError) as caught:
convert_file_document_to_url_document({"type": "file", "file": File()})
assert caught.value is failure
with pytest.raises(litellm.APIConnectionError, match="file property failed") as caught:
call_native_ocr(ocr_server, document={"type": "file", "file": File()})
assert caught.value.__context__ is failure
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_native_file_preparation_preserves_reader_exception(
ocr_server: RecordingServer, asynchronous: bool
) -> None:
ocr_server.expected_requests = 0
failure: Final = RuntimeError("reader failed")
class Reader:
def read(self) -> bytes:
raise failure
document: Final = {"type": "file", "file": Reader()}
with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr(
ocr_server, document=document
)
assert caught.value.__context__ is failure
def test_native_file_preparation_rejects_unsupported_reader_results(ocr_server: RecordingServer) -> None:
ocr_server.expected_requests = 0
class Reader:
def read(self) -> int:
return 1
with pytest.raises(litellm.APIConnectionError, match="bytes or str") as caught:
call_native_ocr(ocr_server, document={"type": "file", "file": Reader()})
assert isinstance(caught.value.__context__, TypeError)
@pytest.mark.parametrize("kind", ["bytes", "path", "reader"])
def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Path) -> None:
from litellm.ocr.input import FileDocument, convert_file_document_to_url_document, get_max_file_bytes
limit: Final = get_max_file_bytes()
def test_native_file_preparation_rejects_oversized_input(ocr_server: RecordingServer, kind: str, tmp_path: Path) -> None:
ocr_server.expected_requests = 0
limit: Final = 50 * 1024 * 1024
path: Final = tmp_path / "large.pdf"
with path.open("wb") as stream:
stream.truncate(limit + 1)
@ -573,53 +606,25 @@ def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Pa
def read(self) -> bytes:
return b"a" * (limit + 1)
document: Final[FileDocument] = {
document: Final = {
"type": "file",
"file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1),
}
with pytest.raises(ValueError, match="exceeds the size limit"):
convert_file_document_to_url_document(document)
with pytest.raises(litellm.BadRequestError, match="exceeds the size limit"):
call_native_ocr(ocr_server, document=document)
@pytest.mark.parametrize("kind", ["str", "path", "reader"])
def test_native_upload_binding_rejects_filesystem_inputs(kind: str, tmp_path: Path) -> None:
def test_native_file_preparation_reports_missing_paths(ocr_server: RecordingServer, tmp_path: Path) -> None:
ocr_server.expected_requests = 0
missing: Final = tmp_path / "missing.pdf"
with pytest.raises(litellm.APIConnectionError, match=f"File not found: {missing}") as caught:
call_native_ocr(ocr_server, document={"type": "file", "file": missing})
assert isinstance(caught.value.__context__, FileNotFoundError)
def test_native_file_preparation_rejects_empty_readers(ocr_server: RecordingServer) -> None:
from io import BytesIO
from typing import cast # noqa: TID251 # deliberately invalid inputs exercise the native runtime boundary
from litellm.ocr.input import convert_upload_to_url_document
path: Final = tmp_path / "secret.pdf"
path.write_bytes(b"server secret")
source: Final = str(path) if kind == "str" else path if kind == "path" else BytesIO(b"abc")
with pytest.raises(TypeError):
convert_upload_to_url_document(cast(bytes, source), "document.pdf", None)
@pytest.mark.parametrize("extra_bytes", [0, 1])
def test_native_upload_enforces_file_size_limit(extra_bytes: int) -> None:
import base64
from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes
content: Final = b"a" * (get_max_file_bytes() + extra_bytes)
if extra_bytes:
with pytest.raises(ValueError, match="exceeds the size limit"):
convert_upload_to_url_document(content, "scan.pdf", None)
return
document: Final = convert_upload_to_url_document(content, "scan.pdf", None)
assert document["type"] == "document_url"
assert base64.b64decode(document["document_url"].split(",", 1)[1]) == content
def test_native_file_preparation_preserves_reader_exception() -> None:
from litellm.ocr.input import convert_file_document_to_url_document
failure: Final = RuntimeError("reader failed")
class Reader:
def read(self) -> bytes:
raise failure
with pytest.raises(RuntimeError) as caught:
convert_file_document_to_url_document({"type": "file", "file": Reader()})
assert caught.value is failure
ocr_server.expected_requests = 0
with pytest.raises(litellm.BadRequestError, match="File is empty"):
call_native_ocr(ocr_server, document={"type": "file", "file": BytesIO(b"")})