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

refactor(ocr): fold execution into shared lifecycle

Browse source
This commit is contained in:
Yujong Lee 2026-09-07 21:38:57 -07:00
parent f1a5a6f47a
commit 0a53a00edd
20 changed files with 824 additions and 607 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -1,6 +1,6 @@
use litellm_core::Error;
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
use litellm_core::lifecycle::CallLifecycle;
use litellm_core::lifecycle::{CallLifecycle, CallLifecycleRequest, SystemClock};
use serde_json::Value;
mod hooks;
@ -14,9 +14,18 @@ use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
let PreparedAudioTranscriptionCall { request, hooks } =
prepare_audio_transcription_call(request);
CallLifecycle::default()
.run_request_result(request, &hooks, execute_audio_transcription_provider_call)
let context = request.lifecycle_context();
CallLifecycle
.run(
context,
request,
&hooks,
&hooks,
&SystemClock,
execute_audio_transcription_provider_call,
)
.await
.into_result()
}
#[cfg(test)]

44
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
View file

@ -3,7 +3,7 @@ use std::time::Duration;
use futures_util::{Sink, Stream};
use litellm_core::Error;
use litellm_core::lifecycle::{CallLifecycle, CallLifecycleContext};
use litellm_core::lifecycle::{CallLifecycle, CallLifecycleContext, SystemClock};
use litellm_core::responses::instrumentation::{
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
ResponsesWsMetadata,
@ -56,23 +56,31 @@ where
));
let observer_instrumentation = Arc::clone(&instrumentation);
let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
let result = CallLifecycle::default()
.run_result(context, (), instrumentation.as_ref(), |_| async move {
crate::io::responses_ws::async_responses_websocket(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
first_frame,
idle_timeout,
move |event| {
observer_instrumentation.observe(event);
},
client_in,
client_out,
)
.await
})
.await;
let result = CallLifecycle
.run(
context,
(),
instrumentation.as_ref(),
instrumentation.as_ref(),
&SystemClock,
|_| async move {
crate::io::responses_ws::async_responses_websocket(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
first_frame,
idle_timeout,
move |event| {
observer_instrumentation.observe(event);
},
client_in,
client_out,
)
.await
},
)
.await
.into_result();
let outcome = instrumentation.take_or_build_outcome(result.is_ok());
dispatch_outcome(loggers, outcome).await;
result

79
litellm-rust/crates/core/src/lifecycle/execution.rs
View file

@ -121,28 +121,12 @@ impl Clock for SystemClock {
pub struct CallLifecycle;
impl CallLifecycle {
pub async fn run<InitialReq, ProviderReq, Resp, Policy, ProviderCall, ProviderFuture>(
&self,
context: CallLifecycleContext,
request: InitialReq,
policy: &Policy,
provider_call: ProviderCall,
) -> ExecutedCall<Resp, Error>
where
Resp: Serialize,
Policy: RequestPolicy<InitialReq, ProviderReq> + TerminalDispatcher,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
self.run_with_clock(context, request, policy, &SystemClock, provider_call)
.await
}
pub async fn run_with_clock<
pub async fn run<
InitialReq,
ProviderReq,
Resp,
Policy,
Dispatcher,
ClockImpl,
ProviderCall,
ProviderFuture,
@ -151,12 +135,14 @@ impl CallLifecycle {
context: CallLifecycleContext,
request: InitialReq,
policy: &Policy,
dispatcher: &Dispatcher,
clock: &ClockImpl,
provider_call: ProviderCall,
) -> ExecutedCall<Resp, Error>
where
Resp: Serialize,
Policy: RequestPolicy<InitialReq, ProviderReq> + TerminalDispatcher,
Policy: RequestPolicy<InitialReq, ProviderReq>,
Dispatcher: TerminalDispatcher,
ClockImpl: Clock,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
@ -165,13 +151,13 @@ impl CallLifecycle {
let request = match policy.async_pre_call_hook(&context, request).await {
ActionResult::Continue(request) | ActionResult::Replace(request) => request,
ActionResult::Reject(error) => {
return failure(policy, clock, &context, error, start_time).await;
return failure(dispatcher, clock, &context, error, start_time).await;
}
};
let provider_request = match policy.async_during_call_hook(&context, request).await {
ActionResult::Continue(request) | ActionResult::Replace(request) => request,
ActionResult::Reject(error) => {
return failure(policy, clock, &context, error, start_time).await;
return failure(dispatcher, clock, &context, error, start_time).await;
}
};
match provider_call(provider_request).await {
@ -181,55 +167,12 @@ impl CallLifecycle {
TerminalClassification::Success,
serde_json::to_value(&response).unwrap_or(Value::Null),
);
let _ = policy.dispatch(&terminal).await;
let _ = dispatcher.dispatch(&terminal).await;
ExecutedCall::Success { response, terminal }
}
Err(error) => failure(policy, clock, &context, error, start_time).await,
Err(error) => failure(dispatcher, clock, &context, error, start_time).await,
}
}
pub async fn run_result<InitialReq, ProviderReq, Resp, Policy, ProviderCall, ProviderFuture>(
&self,
context: CallLifecycleContext,
request: InitialReq,
policy: &Policy,
provider_call: ProviderCall,
) -> Result<Resp, Error>
where
Resp: Serialize,
Policy: RequestPolicy<InitialReq, ProviderReq> + TerminalDispatcher,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
self.run(context, request, policy, provider_call)
.await
.into_result()
}
pub async fn run_request_result<
InitialReq,
ProviderReq,
Resp,
Policy,
ProviderCall,
ProviderFuture,
>(
&self,
request: InitialReq,
policy: &Policy,
provider_call: ProviderCall,
) -> Result<Resp, Error>
where
InitialReq: CallLifecycleRequest,
Resp: Serialize,
Policy: RequestPolicy<InitialReq, ProviderReq> + TerminalDispatcher,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
let context = request.lifecycle_context();
self.run_result(context, request, policy, provider_call)
.await
}
}
async fn failure<R, Policy, ClockImpl>(
@ -350,6 +293,8 @@ mod tests {
CallLifecycleContext::new("ocr", "model", "provider", "call-1"),
"request".to_string(),
&policy,
&policy,
&SystemClock,
|request| async move {
assert_eq!(request, "request:pre:during");
Ok(request)
@ -375,6 +320,8 @@ mod tests {
CallLifecycleContext::new("ocr", "model", "provider", "call-2"),
"request".to_string(),
&policy,
&policy,
&SystemClock,
|request| async move { Ok(request) },
)
.await;

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

@ -1,5 +1,5 @@
use crate::Error;
use crate::ocr::{OcrRequest, prepare};
use crate::ocr::{OcrAdmissionRequest, prepare};
use super::{
ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, LifecycleRoute, Outcome,
@ -89,7 +89,10 @@ pub struct OcrRoute;
pub type Lifecycle = super::Lifecycle<OcrRoute>;
impl Lifecycle {
pub fn new(admission: &OcrRequest, options: Options) -> Result<NativeOutcome<Self>, Error> {
pub fn new(
admission: &OcrAdmissionRequest,
options: Options,
) -> Result<NativeOutcome<Self>, Error> {
<super::Lifecycle<OcrRoute>>::admit(admission, options).map(|admission| match admission {
Ok(lifecycle) => NativeOutcome::Completed(lifecycle),
Err(decline) => NativeOutcome::Declined(decline),
@ -102,7 +105,7 @@ impl Lifecycle {
}
impl LifecycleRoute for OcrRoute {
type Admission = OcrRequest;
type Admission = OcrAdmissionRequest;
type Options = Options;
type Context = Observations;
type Operation = Operation;
@ -262,3 +265,312 @@ fn generate_call_id() -> String {
&hex[20..]
)
}
#[cfg(test)]
mod tests {
use std::rc::Rc;
use super::*;
use crate::ocr::types::OcrDocument;
fn request() -> OcrAdmissionRequest {
OcrAdmissionRequest {
model: "mistral/requested-model".into(),
custom_llm_provider: None,
api_key: Some("test-key".into()),
api_base: Some("https://example.test".into()),
extra_headers: vec![],
timeout_seconds: 2.0,
request_format: None,
document: OcrDocument::DocumentUrl {
document_url: "https://example.test/doc.pdf".into(),
},
azure_ad_token: None,
vertex_project: None,
vertex_location: None,
stream: false,
}
}
fn machine(asynchronous: bool) -> Lifecycle {
let NativeOutcome::Completed(supplied) = Lifecycle::new(
&request(),
Options {
asynchronous,
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
supplied
}
fn observed() -> Observations {
Observations {
logger_available: true,
has_fallbacks: false,
}
}
fn reach(machine: &mut Lifecycle, operation: Operation) {
for _ in 0..12 {
if machine.operation() == operation {
return;
}
machine.advance(Outcome::Success, observed()).unwrap();
}
panic!("operation not reached: {operation:?}")
}
#[test]
fn success_sequences_and_completion_are_core_selected() {
use Operation::*;
for (asynchronous, expected) in [
(false, vec![Setup, Prepare, Send, SyncSuccess, Restore]),
(
true,
vec![
Setup,
DeploymentPre,
Prepare,
Send,
DeploymentSuccess,
AsyncSuccess,
SyncSuccessIfNeeded,
Restore,
],
),
] {
let mut machine = machine(asynchronous);
for operation in expected {
assert_eq!(machine.operation(), operation);
assert_eq!(
machine.advance(Outcome::Success, observed()).unwrap().error,
ErrorDisposition::Preserve
);
}
assert_eq!(machine.operation(), Complete(Outcome::Success));
assert!(machine.advance(Outcome::Success, observed()).is_err());
}
}
#[test]
fn failures_and_cancellation_transition_without_recursion() {
use Operation::*;
for asynchronous in [false, true] {
let stages = if asynchronous {
vec![
Setup,
DeploymentPre,
Prepare,
Send,
DeploymentSuccess,
AsyncSuccess,
SyncSuccessIfNeeded,
]
} else {
vec![Setup, Prepare, Send, SyncSuccess]
};
for stage in stages {
for outcome in [Outcome::Failure, Outcome::Abort] {
let mut machine = machine(asynchronous);
reach(&mut machine, stage);
let transition = machine.advance(outcome, observed()).unwrap();
assert_eq!(transition.error, ErrorDisposition::Replace);
let expected = if outcome == Outcome::Abort {
Restore
} else if asynchronous && matches!(stage, Prepare | Send) {
DeploymentFailure
} else {
SyncFailure
};
assert_eq!(transition.operation, expected, "{stage:?}, {outcome:?}");
if expected == DeploymentFailure {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
SyncFailure
);
}
if outcome == Outcome::Failure {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
if asynchronous { AsyncFailure } else { Restore }
);
if asynchronous {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Restore
);
}
}
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Complete(outcome)
);
}
}
}
}
#[test]
fn deployment_failure_observer_preserves_the_original_error() {
for observer_outcome in [Outcome::Success, Outcome::Failure, Outcome::Abort] {
let mut machine = machine(true);
reach(&mut machine, Operation::Send);
let original = Rc::new("original opaque error");
let mut retained = Rc::clone(&original);
assert_eq!(
machine
.advance(Outcome::Failure, observed())
.unwrap()
.operation,
Operation::DeploymentFailure
);
let transition = machine.advance(observer_outcome, observed()).unwrap();
if transition.error == ErrorDisposition::Replace {
retained = Rc::new("observer error");
}
assert!(Rc::ptr_eq(&original, &retained));
assert_eq!(transition.operation, Operation::SyncFailure);
}
}
#[test]
fn logger_availability_internal_calls_and_fallbacks_control_logging_only() {
let mut failed_setup = machine(true);
assert_eq!(
failed_setup
.advance(Outcome::Failure, Observations::default())
.unwrap()
.operation,
Operation::Restore
);
for internal_call in [false, true] {
for has_fallbacks in [false, true] {
let NativeOutcome::Completed(mut machine) = Lifecycle::new(
&request(),
Options {
asynchronous: true,
internal_call,
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
reach(&mut machine, Operation::DeploymentSuccess);
assert_eq!(
machine
.advance(
Outcome::Success,
Observations {
has_fallbacks,
..observed()
},
)
.unwrap()
.operation,
if internal_call || has_fallbacks {
Operation::SyncSuccessIfNeeded
} else {
Operation::AsyncSuccess
}
);
}
}
}
#[test]
fn admission_declines_unsupported_requests_without_effects() {
for request in [
OcrAdmissionRequest {
document: OcrDocument::File,
..request()
},
OcrAdmissionRequest {
document: OcrDocument::Unsupported,
..request()
},
OcrAdmissionRequest {
stream: true,
..request()
},
OcrAdmissionRequest {
request_format: Some("native".into()),
..request()
},
OcrAdmissionRequest {
model: "openai/model".into(),
..request()
},
] {
assert!(matches!(
Lifecycle::new(&request, Options::default()),
Ok(NativeOutcome::Declined(_))
));
}
assert!(matches!(
Lifecycle::new(
&request(),
Options {
credential_method: CredentialMethod::Acquisition,
..Options::default()
}
),
Ok(NativeOutcome::Declined(_))
));
assert!(matches!(
Lifecycle::new(
&OcrAdmissionRequest {
timeout_seconds: f64::NAN,
..request()
},
Options::default()
),
Err(Error::InvalidRequest(_))
));
}
#[test]
fn identity_preserves_supplied_provenance_and_generates_uuid_v4() {
let NativeOutcome::Completed(supplied) = Lifecycle::new(
&request(),
Options {
call_id: Some("logical-call".into()),
trace_id: Some("trace".into()),
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
assert_eq!(
supplied.identity(),
&Identity {
requested_model: "mistral/requested-model".into(),
call_id: "logical-call".into(),
trace_id: Some("trace".into()),
generated_call_id: false,
}
);
let first = machine(false);
let second = machine(false);
assert!(first.identity().generated_call_id);
assert_eq!(first.identity().call_id.len(), 36);
assert_eq!(&first.identity().call_id[14..15], "4");
assert_ne!(first.identity().call_id, second.identity().call_id);
}
}

11
litellm-rust/crates/core/src/messages/lifecycle.rs
View file

@ -238,9 +238,14 @@ pub async fn messages<S: MessagesServices>(
context: CallLifecycleContext,
) -> ExecutedCall<AnthropicMessagesResponse, Error> {
CallLifecycle
.run_with_clock(context, request, services, services, |request| async move {
execute_messages_provider_call(request).await
})
.run(
context,
request,
services,
services,
services,
|request| async move { execute_messages_provider_call(request).await },
)
.await
}

18
litellm-rust/crates/core/src/ocr/common_utils.rs
View file

@ -1,19 +1,19 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use crate::error::Error;
use crate::ocr::transformation::OcrProviderConfig;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::error::Error;
use litellm_core::ocr::transformation::OcrProviderConfig;
use reqwest::Url;
use serde_json::{Map, Value};
use litellm_core::providers::azure_ai::ocr::transformation as azure_ai;
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use crate::providers::azure_ai::ocr::transformation as azure_ai;
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use crate::providers::reducto::ocr::transformation as reducto;
use crate::providers::vertex_ai::ocr::transformation as vertex_ai;
use crate::client::http_client;
use super::client::http_client;
const ERROR_BODY_MAX_CHARS: usize = 256;
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
@ -55,7 +55,7 @@ pub(super) fn string_headers(
.ok_or_else(|| {
Error::InvalidRequest(format!(
"OCR extra_headers.{key} must be a string, got {}",
litellm_core::error::json_type_name(&value)
crate::error::json_type_name(&value)
))
})
})
@ -381,7 +381,7 @@ pub(super) async fn poll_document_intelligence(
#[cfg(test)]
mod tests {
use litellm_core::ocr::transformation::OcrResponseHandling;
use crate::ocr::transformation::OcrResponseHandling;
use serde_json::json;
use super::*;

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

@ -5,15 +5,15 @@ use serde_json::Value;
use super::client::http_client;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::hooks::OcrLifecycleHooks;
use super::hooks::OcrRequestPolicy;
use super::runtime_types::PreparedOcrRequest;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
policy: &OcrRequestPolicy,
) -> Result<Value, Error> {
let request = hooks.prepare_provider_request(request).await?;
let request = policy.prepare_provider_request(request).await?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);

49
litellm-rust/crates/core/src/ocr/hooks.rs
View file

@ -1,7 +1,5 @@
use crate::error::Error;
use crate::lifecycle::{
ActionResult, CallLifecycleContext, RequestPolicy, TerminalDispatcher, TerminalRecord,
};
use crate::lifecycle::{ActionResult, CallLifecycleContext, RequestPolicy};
use crate::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
@ -15,25 +13,22 @@ use super::runtime_types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{CallType, CustomLoggerRunner, LogFuture};
use crate::integrations::custom_logger::CallType;
use crate::integrations::types::RequestMetadata;
pub(crate) struct OcrLifecycleHooks {
logger_runner: CustomLoggerRunner,
pub(crate) struct OcrRequestPolicy {
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = ActionResult<T, Error>> + Send + 'a>>;
impl OcrLifecycleHooks {
impl OcrRequestPolicy {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
@ -153,7 +148,6 @@ impl OcrLifecycleHooks {
.map_err(guardrail_error_to_core_error)?;
parse_ocr_during_call_guardrail_request(guardrail_request)
}
}
async fn upload_reducto_document(
@ -213,7 +207,7 @@ async fn upload_reducto_document(
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl RequestPolicy<PreparedOcrRequest, PreparedOcrRequest> for OcrLifecycleHooks {
impl RequestPolicy<PreparedOcrRequest, PreparedOcrRequest> for OcrRequestPolicy {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
@ -239,23 +233,6 @@ impl RequestPolicy<PreparedOcrRequest, PreparedOcrRequest> for OcrLifecycleHooks
}
}
impl TerminalDispatcher for OcrLifecycleHooks {
fn dispatch<'a>(&'a self, terminal: &'a TerminalRecord) -> LogFuture<'a> {
let mut terminal = terminal.clone();
terminal.cost_inputs.metadata = request_metadata(&self.request_metadata);
Box::pin(async move { self.logger_runner.dispatch(&terminal).await })
}
}
fn request_metadata(metadata: &RequestMetadata) -> crate::integrations::types::StandardLoggingMetadata {
crate::integrations::types::StandardLoggingMetadata {
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
..Default::default()
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Ocr,
@ -304,19 +281,3 @@ fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

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

@ -1,409 +0,0 @@
pub use crate::lifecycle::ocr::*;
pub use crate::lifecycle::{ErrorDisposition, Outcome};
#[cfg(test)]
use crate::Error;
#[cfg(test)]
use crate::ocr::{OcrRequest, prepare};
#[cfg(test)]
mod tests {
use std::rc::Rc;
use super::*;
use crate::ocr::types::OcrDocument;
fn request() -> OcrRequest {
OcrRequest {
model: "mistral/requested-model".into(),
custom_llm_provider: None,
api_key: Some("test-key".into()),
api_base: Some("https://example.test".into()),
extra_headers: vec![],
timeout_seconds: 2.0,
request_format: None,
document: OcrDocument::DocumentUrl {
document_url: "https://example.test/doc.pdf".into(),
},
azure_ad_token: None,
vertex_project: None,
vertex_location: None,
stream: false,
}
}
fn machine(asynchronous: bool) -> Lifecycle {
let NativeOutcome::Completed(machine) = Lifecycle::new(
&request(),
Options {
asynchronous,
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
machine
}
fn observed() -> Observations {
Observations {
logger_available: true,
has_fallbacks: false,
}
}
fn reach(machine: &mut Lifecycle, operation: Operation) {
for _ in 0..12 {
if machine.operation() == operation {
return;
}
machine.advance(Outcome::Success, observed()).unwrap();
}
panic!("operation not reached: {operation:?}")
}
#[test]
fn success_sequences_and_completion_are_core_selected() {
use Operation::*;
for (asynchronous, expected) in [
(false, vec![Setup, Prepare, Send, SyncSuccess, Restore]),
(
true,
vec![
Setup,
DeploymentPre,
Prepare,
Send,
DeploymentSuccess,
AsyncSuccess,
SyncSuccessIfNeeded,
Restore,
],
),
] {
let mut machine = machine(asynchronous);
for operation in expected {
assert_eq!(machine.operation(), operation);
assert_eq!(
machine.advance(Outcome::Success, observed()).unwrap().error,
ErrorDisposition::Preserve
);
}
assert_eq!(machine.operation(), Complete(Outcome::Success));
assert!(machine.advance(Outcome::Success, observed()).is_err());
}
}
#[test]
fn ordinary_errors_and_cancellation_at_every_execution_stage() {
use Operation::*;
for asynchronous in [false, true] {
let stages = if asynchronous {
vec![
Setup,
DeploymentPre,
Prepare,
Send,
DeploymentSuccess,
AsyncSuccess,
SyncSuccessIfNeeded,
]
} else {
vec![Setup, Prepare, Send, SyncSuccess]
};
for stage in stages {
for outcome in [Outcome::Failure, Outcome::Abort] {
let mut machine = machine(asynchronous);
reach(&mut machine, stage);
let transition = machine.advance(outcome, observed()).unwrap();
assert_eq!(transition.error, ErrorDisposition::Replace);
let expected = if outcome == Outcome::Abort {
Restore
} else if asynchronous && matches!(stage, Prepare | Send) {
DeploymentFailure
} else {
SyncFailure
};
assert_eq!(transition.operation, expected, "{stage:?}, {outcome:?}");
if expected == DeploymentFailure {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
SyncFailure
);
}
if outcome == Outcome::Failure {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
if asynchronous { AsyncFailure } else { Restore }
);
if asynchronous {
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Restore
);
}
}
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Complete(outcome)
);
}
}
}
}
#[test]
fn deployment_observer_preserves_opaque_original_error_even_on_abort() {
for observer_outcome in [Outcome::Success, Outcome::Failure, Outcome::Abort] {
let mut machine = machine(true);
reach(&mut machine, Operation::Send);
let original = Rc::new("original opaque error");
let mut retained = Rc::clone(&original);
assert_eq!(
machine
.advance(Outcome::Failure, observed())
.unwrap()
.operation,
Operation::DeploymentFailure
);
let transition = machine.advance(observer_outcome, observed()).unwrap();
if transition.error == ErrorDisposition::Replace {
retained = Rc::new("observer error");
}
assert!(Rc::ptr_eq(&original, &retained));
assert_eq!(transition.operation, Operation::SyncFailure);
reach(&mut machine, Operation::Restore);
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Operation::Complete(Outcome::Failure)
);
}
}
#[test]
fn failure_handlers_and_restore_propagate_their_own_errors_without_recursion() {
for stage in [
Operation::SyncFailure,
Operation::AsyncFailure,
Operation::Restore,
] {
for outcome in [Outcome::Failure, Outcome::Abort] {
let mut machine = machine(true);
machine.advance(Outcome::Failure, observed()).unwrap();
reach(&mut machine, stage);
let transition = machine.advance(outcome, observed()).unwrap();
assert_eq!(transition.error, ErrorDisposition::Replace);
if stage != Operation::Restore {
assert_eq!(transition.operation, Operation::Restore);
machine.advance(Outcome::Success, observed()).unwrap();
}
assert_eq!(machine.operation(), Operation::Complete(outcome));
}
}
}
#[test]
fn logger_availability_internal_calls_and_fallbacks_control_logging_only() {
let mut failed_setup = machine(true);
assert_eq!(
failed_setup
.advance(Outcome::Failure, Observations::default())
.unwrap()
.operation,
Operation::Restore
);
for internal_call in [false, true] {
for has_fallbacks in [false, true] {
let NativeOutcome::Completed(mut machine) = Lifecycle::new(
&request(),
Options {
asynchronous: true,
internal_call,
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
reach(&mut machine, Operation::DeploymentSuccess);
let next = machine
.advance(
Outcome::Success,
Observations {
has_fallbacks,
..observed()
},
)
.unwrap();
assert_eq!(
next.operation,
if internal_call || has_fallbacks {
Operation::SyncSuccessIfNeeded
} else {
Operation::AsyncSuccess
}
);
}
}
let NativeOutcome::Completed(mut internal) = Lifecycle::new(
&request(),
Options {
asynchronous: true,
internal_call: true,
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
reach(&mut internal, Operation::Prepare);
assert_eq!(
internal
.advance(Outcome::Failure, observed())
.unwrap()
.operation,
Operation::DeploymentFailure
);
assert_eq!(
internal
.advance(Outcome::Abort, observed())
.unwrap()
.operation,
Operation::Restore
);
}
#[test]
fn decline_is_admission_only_and_file_is_an_inert_descriptor() {
for request in [
OcrRequest {
document: OcrDocument::File,
..request()
},
OcrRequest {
document: OcrDocument::Unsupported,
..request()
},
OcrRequest {
stream: true,
..request()
},
OcrRequest {
request_format: Some("native".into()),
..request()
},
OcrRequest {
model: "openai/model".into(),
..request()
},
] {
assert!(matches!(
Lifecycle::new(&request, Options::default()),
Ok(NativeOutcome::Declined(_))
));
}
assert!(matches!(
Lifecycle::new(
&request(),
Options {
credential_method: CredentialMethod::Acquisition,
..Options::default()
}
),
Ok(NativeOutcome::Declined(_))
));
assert!(matches!(
Lifecycle::new(
&OcrRequest {
timeout_seconds: f64::NAN,
..request()
},
Options::default()
),
Err(Error::InvalidRequest(_))
));
let mut machine = machine(true);
reach(&mut machine, Operation::Prepare);
assert!(matches!(
prepare::prepare(OcrRequest {
document: OcrDocument::File,
..request()
}),
Err(Error::Unsupported(_))
));
assert_eq!(
machine
.advance(Outcome::Failure, observed())
.unwrap()
.operation,
Operation::DeploymentFailure
);
reach(&mut machine, Operation::Restore);
assert_eq!(
machine
.advance(Outcome::Success, observed())
.unwrap()
.operation,
Operation::Complete(Outcome::Failure)
);
}
#[test]
fn identity_keeps_supplied_provenance_across_hook_replacement_and_sdk_attempts() {
for _ in 0..2 {
let NativeOutcome::Completed(mut machine) = Lifecycle::new(
&request(),
Options {
asynchronous: true,
call_id: Some("logical-call".into()),
trace_id: Some("trace".into()),
..Options::default()
},
)
.unwrap() else {
panic!("expected admission")
};
reach(&mut machine, Operation::Prepare);
let prepared = prepare::prepare(OcrRequest {
model: "mistral/replacement".into(),
..request()
})
.unwrap();
assert_eq!(prepared.model, "replacement");
assert_eq!(
machine.identity(),
&Identity {
requested_model: "mistral/requested-model".into(),
call_id: "logical-call".into(),
trace_id: Some("trace".into()),
generated_call_id: false,
}
);
reach(&mut machine, Operation::Restore);
assert_eq!(machine.identity().call_id, "logical-call");
}
let first = machine(false);
let second = machine(false);
assert!(first.identity().generated_call_id);
assert_eq!(first.identity().call_id.len(), 36);
assert_eq!(&first.identity().call_id[14..15], "4");
assert_ne!(first.identity().call_id, second.identity().call_id);
assert_eq!(first.identity().trace_id, None);
}
}

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

@ -1,5 +1,9 @@
pub mod lifecycle;
mod client;
mod common_utils;
mod handler;
mod hooks;
pub mod prepare;
mod runtime_types;
pub mod transformation;
pub mod types;
@ -9,19 +13,167 @@ use crate::Error;
use crate::error::json_type_name;
use crate::http_utils::{buffered_post, has_header};
pub use runtime_types::{OcrRequest, OcrRouteRequest};
pub use types::{OcrAdmissionRequest, OcrResponseData, PreparedOcr, PreparedOcrCall};
use types::{OcrDocument, OcrDocumentProjection};
pub use types::{OcrRequest, OcrResponseData, PreparedOcr};
pub fn terminal_callbacks(asynchronous: bool, success: bool) -> &'static [&'static str] {
match (asynchronous, success) {
(false, true) => &["sync_success"],
(true, true) => &["async_success", "sync_success_if_needed"],
(false, false) => &["sync_failure"],
(true, false) => &["sync_failure", "async_failure"],
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
use crate::lifecycle::{
CallLifecycle, CallLifecycleContext, Clock, ExecutedCall, TerminalDispatcher,
};
use hooks::OcrRequestPolicy;
use runtime_types::PreparedOcrRequest;
pub trait OcrServices: TerminalDispatcher + Clock {}
impl<T> OcrServices for T where T: TerminalDispatcher + Clock {}
pub struct DefaultOcrServices {
dispatcher: CustomLoggerRunner,
}
pub struct NoopOcrServices;
impl Default for NoopOcrServices {
fn default() -> Self {
Self
}
}
pub async fn ocr(
impl DefaultOcrServices {
pub fn new(request: &OcrRequest<'_>) -> Self {
Self {
dispatcher: CustomLoggerRunner::new(request.callbacks.clone()),
}
}
}
impl Clock for DefaultOcrServices {
fn now(&self) -> f64 {
crate::lifecycle::SystemClock.now()
}
}
impl TerminalDispatcher for DefaultOcrServices {
fn dispatch<'a>(
&'a self,
terminal: &'a crate::lifecycle::TerminalRecord,
) -> crate::integrations::custom_logger::LogFuture<'a> {
self.dispatcher.dispatch(terminal)
}
}
impl Clock for NoopOcrServices {
fn now(&self) -> f64 {
crate::lifecycle::SystemClock.now()
}
}
impl TerminalDispatcher for NoopOcrServices {
fn dispatch<'a>(
&'a self,
_: &'a crate::lifecycle::TerminalRecord,
) -> crate::integrations::custom_logger::LogFuture<'a> {
Box::pin(async { Ok(()) })
}
}
pub async fn ocr<S: OcrServices>(
services: &S,
request: OcrRouteRequest<'_>,
_options: crate::lifecycle::ocr::Options,
context: CallLifecycleContext,
) -> ExecutedCall<Value, Error> {
let OcrRouteRequest::Native(request) = request else {
let OcrRouteRequest::Prepared(request) = request else {
unreachable!()
};
return CallLifecycle
.run(
context,
request,
&PreparedOcrPolicy,
services,
services,
|request| async move {
send_prepared(request.prepared, request.headers, request.body)
.await
.map(OcrResponseData::into_json)
},
)
.await;
};
let policy = OcrRequestPolicy::new(
CustomGuardrailRunner::new(request.guardrails.clone()),
request.request_metadata.clone(),
);
let provider = crate::routing_utils::provider::get_custom_llm_provider(
request.model,
request.custom_llm_provider,
);
let config = provider
.as_ref()
.ok_or_else(|| Error::InvalidProvider("unable to resolve OCR provider".into()))
.and_then(|provider| {
common_utils::ocr_provider_config(provider.custom_llm_provider, provider.model)
.ok_or_else(|| Error::InvalidProvider("unsupported OCR provider".into()))
});
let provider_model = provider.as_ref().map_or(request.model, |value| value.model);
let provider_name = provider
.as_ref()
.map_or(request.custom_llm_provider.unwrap_or(""), |value| {
value.custom_llm_provider
});
let prepared = PreparedOcrRequest {
config,
model: provider_model.to_string(),
custom_llm_provider: provider_name.to_string(),
litellm_call_id: context.litellm_call_id.clone(),
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
timeout: request.timeout,
};
CallLifecycle
.run(context, prepared, &policy, services, services, |request| {
handler::execute_ocr_provider_call(request, &policy)
})
.await
}
struct PreparedOcrPolicy;
impl crate::lifecycle::RequestPolicy<PreparedOcrCall, PreparedOcrCall> for PreparedOcrPolicy {
type PreCallFuture<'a>
= std::future::Ready<crate::lifecycle::ActionResult<PreparedOcrCall, Error>>
where
Self: 'a;
type DuringCallFuture<'a>
= std::future::Ready<crate::lifecycle::ActionResult<PreparedOcrCall, Error>>
where
Self: 'a;
fn async_pre_call_hook<'a>(
&'a self,
_: &'a CallLifecycleContext,
request: PreparedOcrCall,
) -> Self::PreCallFuture<'a> {
std::future::ready(crate::lifecycle::ActionResult::Continue(request))
}
fn async_during_call_hook<'a>(
&'a self,
_: &'a CallLifecycleContext,
request: PreparedOcrCall,
) -> Self::DuringCallFuture<'a> {
std::future::ready(crate::lifecycle::ActionResult::Continue(request))
}
}
pub(crate) async fn send_prepared(
prepared: PreparedOcr,
headers: Vec<(String, String)>,
body: Value,

22
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -9,11 +9,11 @@ use crate::providers::vertex_ai::ocr::transformation as vertex_ai;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{OcrProviderConfig, OcrResponseHandling};
use super::types::OcrRequest;
use super::types::OcrAdmissionRequest;
pub use super::types::PreparedOcr;
fn request_config(
request: &OcrRequest,
request: &OcrAdmissionRequest,
) -> Result<(CustomLlmProvider<'_>, &'static dyn OcrProviderConfig), Error> {
match request.request_format.as_deref() {
None | Some("litellm") => {}
@ -38,12 +38,12 @@ fn request_config(
Ok((provider, config))
}
pub(crate) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> {
pub(crate) fn admission_capabilities(request: &OcrAdmissionRequest) -> Result<(), Error> {
check_admission_capabilities(request, &|key| std::env::var(key).ok())
}
fn check_admission_capabilities(
request: &OcrRequest,
request: &OcrAdmissionRequest,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<(), Error> {
let (provider, config) = request_config(request)?;
@ -81,7 +81,7 @@ fn check_admission_capabilities(
Ok(())
}
pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
pub fn prepare(request: OcrAdmissionRequest) -> Result<PreparedOcr, Error> {
let (provider, config) = request_config(&request)?;
let env_lookup = |key: &str| std::env::var(key).ok();
let headers = config
@ -169,8 +169,8 @@ mod tests {
use super::*;
use crate::ocr::types::OcrDocument;
fn request() -> OcrRequest {
OcrRequest {
fn request() -> OcrAdmissionRequest {
OcrAdmissionRequest {
model: "mistral/mistral-ocr-latest".into(),
custom_llm_provider: None,
api_key: Some("test-key".into()),
@ -207,7 +207,7 @@ mod tests {
#[case] format: Option<&str>,
#[case] expected: &str,
) {
let request = OcrRequest {
let request = OcrAdmissionRequest {
model: model.into(),
request_format: format.map(str::to_owned),
stream: true,
@ -240,7 +240,7 @@ mod tests {
#[test]
fn admission_does_not_validate_credentials_or_resolve_headers() {
let request = OcrRequest {
let request = OcrAdmissionRequest {
api_key: None,
..request()
};
@ -265,7 +265,7 @@ mod tests {
#[case] model: &str,
#[case] env_name: &str,
) {
let request = OcrRequest {
let request = OcrAdmissionRequest {
model: model.into(),
api_key: None,
document: OcrDocument::ImageUrl {
@ -286,7 +286,7 @@ mod tests {
.then(|| "configured".into()))
.is_ok()
);
let request = OcrRequest {
let request = OcrAdmissionRequest {
extra_headers: vec![("AUTHORIZATION".into(), "inert".into())],
..request
};

18
litellm-rust/crates/core/src/ocr/runtime_types.rs
View file

@ -5,6 +5,7 @@ use crate::lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use crate::ocr::transformation::OcrProviderConfig;
use serde_json::{Map, Value};
use super::types::PreparedOcrCall;
use crate::integrations::custom_guardrail::CustomGuardrail;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
@ -24,6 +25,23 @@ pub struct OcrRequest<'a> {
pub litellm_call_id: Option<&'a str>,
}
pub enum OcrRouteRequest<'a> {
Native(OcrRequest<'a>),
Prepared(PreparedOcrCall),
}
impl<'a> From<OcrRequest<'a>> for OcrRouteRequest<'a> {
fn from(request: OcrRequest<'a>) -> Self {
Self::Native(request)
}
}
impl From<PreparedOcrCall> for OcrRouteRequest<'static> {
fn from(request: PreparedOcrCall) -> Self {
Self::Prepared(request)
}
}
pub(crate) struct PreparedOcrRequest {
pub(crate) config: Result<&'static dyn OcrProviderConfig, crate::Error>,
pub(crate) model: String,

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

@ -1,7 +1,7 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
pub struct OcrRequest {
pub struct OcrAdmissionRequest {
pub model: String,
pub custom_llm_provider: Option<String>,
pub api_key: Option<String>,
@ -16,6 +16,12 @@ pub struct OcrRequest {
pub stream: bool,
}
pub struct PreparedOcrCall {
pub prepared: PreparedOcr,
pub headers: Vec<(String, String)>,
pub body: Value,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum OcrDocument {
@ -80,10 +86,14 @@ pub struct OcrRequestData {
pub struct OcrResponseData {
pub pages: Vec<Value>,
pub model: String,
#[serde(default)]
pub document_annotation: Option<Value>,
#[serde(default)]
pub usage_info: Option<Value>,
pub object: String,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
#[serde(default)]
pub provider_native_response: Option<Value>,
}

2
litellm-rust/crates/core/src/responses/instrumentation.rs
View file

@ -330,6 +330,8 @@ mod tests {
),
(),
&instrumentation,
&instrumentation,
&crate::lifecycle::SystemClock,
|_| async { Ok::<(), Error>(()) },
)
.await;

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

@ -5,9 +5,12 @@ use std::thread;
use std::time::Duration;
use litellm_core::Error;
use litellm_core::lifecycle::CallLifecycleContext;
use litellm_core::ocr::prepare::prepare;
use litellm_core::ocr::types::{OcrDocument, OcrDocumentProjection};
use litellm_core::ocr::{OcrRequest, PreparedOcr, ocr};
use litellm_core::ocr::{
NoopOcrServices, OcrAdmissionRequest as OcrRequest, PreparedOcr, PreparedOcrCall,
};
use serde_json::{Value, json};
fn request() -> OcrRequest {
@ -40,6 +43,29 @@ fn body(prepared: &PreparedOcr) -> Value {
Value::Object(body)
}
async fn ocr(
prepared: PreparedOcr,
headers: Vec<(String, String)>,
body: Value,
) -> Result<litellm_core::ocr::OcrResponseData, Error> {
let model = prepared.model.clone();
let provider = prepared.custom_llm_provider.clone();
let response = litellm_core::ocr::ocr(
&NoopOcrServices,
PreparedOcrCall {
prepared,
headers,
body,
}
.into(),
Default::default(),
CallLifecycleContext::new("ocr", model, provider, "test-call"),
)
.await
.into_result()?;
serde_json::from_value(response).map_err(|error| Error::InvalidResponse(error.to_string()))
}
#[test]
fn prepares_provider_template_auth_and_url() {
let prepared = prepare(OcrRequest {

27
litellm-rust/crates/core/tests/ocr_lifecycle.rs
View file

@ -10,9 +10,10 @@ use litellm_core::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use litellm_core::integrations::types::RequestMetadata;
use litellm_core::lifecycle::CallLifecycleContext;
#[cfg(feature = "observability")]
use litellm_core::observability::FunctionTrace;
use litellm_core::ocr::{OcrRequest, ocr};
use litellm_core::ocr::{DefaultOcrServices, OcrRequest};
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
@ -35,6 +36,30 @@ async fn read_http_headers(socket: &mut TcpStream) -> String {
String::from_utf8(request).expect("request is utf8")
}
async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
let services = DefaultOcrServices::new(&request);
let provider = request
.custom_llm_provider
.or_else(|| request.model.split_once('/').map(|(provider, _)| provider))
.unwrap_or("");
let metadata = litellm_core::integrations::types::StandardLoggingMetadata {
user_api_key_hash: request.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: request.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: request.request_metadata.user_api_key_team_id.clone(),
..Default::default()
};
let context = CallLifecycleContext::new(
"ocr",
request.model,
provider,
request.litellm_call_id.unwrap_or(""),
)
.with_metadata(metadata);
litellm_core::ocr::ocr(&services, request.into(), Default::default(), context)
.await
.into_result()
}
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];

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

@ -313,7 +313,7 @@ class Host:
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'amessages')
def terminal(self, action, value):
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, value, self.start, self.end)
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end)
def sync_success(self): return self.terminal('sync_success', self.response)
def async_success(self): return self.terminal('async_success', self.response)

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

@ -3,10 +3,14 @@
//! callbacks; no Python preparation, auth, encoding, or provider transforms run.
use litellm_core::error::Error;
use litellm_core::ocr::lifecycle::{
ErrorDisposition, Lifecycle, NativeOutcome, Observations, Operation, Options, Outcome,
use litellm_core::lifecycle::ocr::{NativeOutcome, Observations, OcrRoute, Operation, Options};
use litellm_core::lifecycle::{
CallLifecycleContext, ErrorDisposition, ExecutedCall, Lifecycle, Outcome, TerminalRecord,
};
use litellm_core::ocr::NoopOcrServices;
use litellm_core::ocr::types::{
OcrAdmissionRequest, OcrDocumentProjection, PreparedOcr, PreparedOcrCall,
};
use litellm_core::ocr::types::{OcrDocumentProjection, OcrRequest, PreparedOcr};
use litellm_core::routing_utils::provider::get_custom_llm_provider;
use litellm_python_interop::{Pythonized, from_py, to_py};
use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError};
@ -26,6 +30,7 @@ struct OcrState {
headers: Option<Py<PyDict>>,
logging: Option<Py<PyAny>>,
prepared: Option<PreparedOcr>,
terminal: Option<TerminalRecord>,
}
#[pymethods]
@ -45,6 +50,7 @@ impl OcrState {
state.body.take(),
state.headers.take(),
state.logging.take(),
state.terminal.take(),
)
};
drop(roots);
@ -123,7 +129,7 @@ fn header_pairs(headers: &Bound<'_, PyDict>) -> PyResult<Vec<(String, String)>>
.collect()
}
fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult<OcrRequest> {
fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult<OcrAdmissionRequest> {
let document = bag
.get_item("document")?
.ok_or_else(|| PyValueError::new_err("OCR requires document"))?
@ -153,7 +159,7 @@ fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult<OcrReques
document_input.set_item(name, value)?;
}
}
Ok(OcrRequest {
Ok(OcrAdmissionRequest {
model,
custom_llm_provider,
api_key: scalar(bag, "api_key")?,
@ -193,7 +199,7 @@ fn request_error_to_pyerr(
#[pyclass]
struct OcrLifecycle {
machine: Lifecycle,
machine: Lifecycle<OcrRoute>,
asynchronous: bool,
}
@ -399,6 +405,7 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResul
headers: Some(headers.unbind()),
logging: Some(logging),
prepared: Some(prepared),
terminal: None,
},
)
}
@ -435,16 +442,47 @@ fn request(py: Python<'_>, state: &Py<OcrState>) -> PyResult<OcrWireRequest> {
fn send(py: Python<'_>, state: Py<OcrState>) -> PyResult<Bound<'_, PyAny>> {
let (prepared, headers, body) = request(py, &state)?;
litellm_python_interop::run_async_py(py, async move {
let _state = state;
let model = prepared.model.clone();
let provider = prepared.custom_llm_provider.clone();
let response = run_async_value(
async move { Ok(litellm_core::ocr::ocr(prepared, headers, body).await) },
core_error_to_pyerr,
let error_model = model.clone();
let error_provider = provider.clone();
let call_id = Python::attach(|py| {
state
.borrow(py)
.arguments
.as_ref()
.and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten())
.unwrap_or_default()
});
let executed = run_async_value(
async move {
let services = NoopOcrServices;
Ok::<_, std::convert::Infallible>(
litellm_core::ocr::ocr(
&services,
PreparedOcrCall {
prepared,
headers,
body,
}
.into(),
Options::default(),
CallLifecycleContext::new("ocr", &model, &provider, call_id),
)
.await,
)
},
|never| match never {},
)
.await?
.map_err(|error| Python::attach(|py| ocr_error_to_pyerr(py, error, &model, &provider)))?;
Ok(Pythonized(response.into_json()))
.await?;
let terminal = executed.terminal().clone();
Python::attach(|py| state.borrow_mut(py).terminal = Some(terminal));
match executed {
ExecutedCall::Success { response, .. } => Ok(Pythonized(response)),
ExecutedCall::Failure { error, .. } => Err(Python::attach(|py| {
ocr_error_to_pyerr(py, error, &error_model, &error_provider)
})),
}
})
}
@ -473,20 +511,58 @@ fn send_sync(py: Python<'_>, state: Py<OcrState>) -> PyResult<Py<PyAny>> {
let (prepared, headers, body) = request(py, &state)?;
let model = prepared.model.clone();
let provider = prepared.custom_llm_provider.clone();
let response = run_sync_value(
let error_model = model.clone();
let error_provider = provider.clone();
let call_id = state
.borrow(py)
.arguments
.as_ref()
.and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten())
.unwrap_or_default();
let executed = run_sync_value(
py,
async move { Ok(litellm_core::ocr::ocr(prepared, headers, body).await) },
core_error_to_pyerr,
)?
.map_err(|error| ocr_error_to_pyerr(py, error, &model, &provider))?;
let fields = to_py(py, &response.into_json())?
.into_bound(py)
.cast_into::<PyDict>()?;
async move {
let services = NoopOcrServices;
Ok::<_, std::convert::Infallible>(
litellm_core::ocr::ocr(
&services,
PreparedOcrCall {
prepared,
headers,
body,
}
.into(),
Options::default(),
CallLifecycleContext::new("ocr", &model, &provider, call_id),
)
.await,
)
},
|never| match never {},
)?;
state.borrow_mut(py).terminal = Some(executed.terminal().clone());
let response = match executed {
ExecutedCall::Success { response, .. } => response,
ExecutedCall::Failure { error, .. } => {
return Err(ocr_error_to_pyerr(py, error, &error_model, &error_provider));
}
};
let fields = to_py(py, &response)?.into_bound(py).cast_into::<PyDict>()?;
let response = finish(py, fields.unbind());
drop(state);
response
}
#[pyfunction]
fn terminal_record(py: Python<'_>, state: Py<OcrState>) -> PyResult<Py<PyAny>> {
let terminal = state
.borrow(py)
.terminal
.clone()
.ok_or_else(|| PyRuntimeError::new_err("OCR terminal record is unavailable"))?;
to_py(py, &terminal)
}
#[pyfunction]
fn ocr(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
driver(py)?.getattr("drive_sync")?.call1((arguments,))
@ -558,7 +634,8 @@ class Host:
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'aocr')
def terminal(self, action, value):
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, value, self.start, self.end)
record = _terminal_record(self.state) if self.state is not None else None
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, record, value, self.start, self.end)
def sync_success(self):
return self.terminal('sync_success', self.response)
@ -600,6 +677,10 @@ class Host:
module.add("_send", wrap_pyfunction!(send, &module)?)?;
module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
module.add(
"_terminal_record",
wrap_pyfunction!(terminal_record, &module)?,
)?;
Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
}
@ -617,6 +698,10 @@ pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
#[cfg(test)]
mod tests {
use super::*;
use litellm_core::integrations::custom_logger::CallbackTiming;
use litellm_core::integrations::types::Usage;
use litellm_core::lifecycle::{RouteProjection, TerminalClassification};
use serde_json::json;
#[test]
#[ignore = "requires the Python SDK and its dependencies on PYTHONPATH"]
@ -730,6 +815,58 @@ mod tests {
});
}
#[test]
fn terminal_record_exports_core_timing() {
Python::initialize();
Python::attach(|py| {
let state = Py::new(
py,
OcrState {
arguments: None,
body: None,
headers: None,
logging: None,
prepared: None,
terminal: Some(TerminalRecord {
call_id: "call-1".into(),
trace_id: None,
attempt: 1,
call_type: "ocr".into(),
model: "model".into(),
provider: "mistral".into(),
timing: CallbackTiming::new(10.25, 12.5),
usage: Usage::default(),
cost_inputs: Default::default(),
classification: TerminalClassification::Success,
projection: RouteProjection::Ocr {
value: json!({"pages": []}),
},
}),
},
)
.unwrap();
let record = terminal_record(py, state).unwrap();
let timing = record.bind(py).get_item("timing").unwrap();
assert_eq!(
timing
.get_item("start_time")
.unwrap()
.extract::<f64>()
.unwrap(),
10.25
);
assert_eq!(
timing
.get_item("end_time")
.unwrap()
.extract::<f64>()
.unwrap(),
12.5
);
});
}
#[test]
fn callback_decline_is_terminal_and_identity_is_reused() {
Python::initialize();

3
litellm/rust_bridge/_lifecycle.py
View file

@ -13,10 +13,11 @@ def invoke_terminal(
action: str,
roots: object,
logger: object,
record: dict[str, object] | None,
value: object,
start_time: datetime,
end_time: datetime,
) -> object:
from litellm.rust_bridge.ocr import invoke_terminal as invoke_ocr_terminal
return invoke_ocr_terminal(action, roots, logger, value, start_time, end_time)
return invoke_ocr_terminal(action, roots, logger, record, value, start_time, end_time)

17
litellm/rust_bridge/ocr.py
View file

@ -191,13 +191,26 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route:
def invoke_terminal(
action: str, roots: object, logger: object, value: object, start_time: datetime, end_time: datetime
action: str,
roots: object,
logger: object,
record: dict[str, object] | None,
value: object,
fallback_start_time: datetime,
fallback_end_time: datetime,
) -> object:
from litellm import utils
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging: Final = cast(Logging, logger) # cast-ok: Rust passes the logger returned by initialize_logging
timing: Final = cast(dict[str, object], record["timing"]) if record is not None else None
start_time: Final = (
datetime.fromtimestamp(cast(float, timing["start_time"])) if timing is not None else fallback_start_time
)
end_time: Final = (
datetime.fromtimestamp(cast(float, timing["end_time"])) if timing is not None else fallback_end_time
)
if action == "sync_success":
def run() -> None:
@ -221,7 +234,7 @@ def invoke_terminal(
return None
if action == "sync_success_if_needed":
if logging._should_run_sync_callbacks_for_async_calls(): # pyright: ignore[reportPrivateUsage] # preserves Logging's async callback policy
return invoke_terminal("sync_success", roots, logger, value, start_time, end_time)
return invoke_terminal("sync_success", roots, logger, record, value, fallback_start_time, fallback_end_time)
return None
exception: Final = cast(Exception, value) # cast-ok: Rust routes terminal failure values as Python exceptions
trace: Final = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))