diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index b39a4b5600a..b0d6c962dcf 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -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 { 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)] diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs index e85fa16135b..9c4067f6da6 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs @@ -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 diff --git a/litellm-rust/crates/core/src/lifecycle/execution.rs b/litellm-rust/crates/core/src/lifecycle/execution.rs index b4d58e8697e..233db6e56a9 100644 --- a/litellm-rust/crates/core/src/lifecycle/execution.rs +++ b/litellm-rust/crates/core/src/lifecycle/execution.rs @@ -121,28 +121,12 @@ impl Clock for SystemClock { pub struct CallLifecycle; impl CallLifecycle { - pub async fn run( - &self, - context: CallLifecycleContext, - request: InitialReq, - policy: &Policy, - provider_call: ProviderCall, - ) -> ExecutedCall - where - Resp: Serialize, - Policy: RequestPolicy + TerminalDispatcher, - ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, - { - 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 where Resp: Serialize, - Policy: RequestPolicy + TerminalDispatcher, + Policy: RequestPolicy, + Dispatcher: TerminalDispatcher, ClockImpl: Clock, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, ProviderFuture: Future>, @@ -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( - &self, - context: CallLifecycleContext, - request: InitialReq, - policy: &Policy, - provider_call: ProviderCall, - ) -> Result - where - Resp: Serialize, - Policy: RequestPolicy + TerminalDispatcher, - ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, - { - 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 - where - InitialReq: CallLifecycleRequest, - Resp: Serialize, - Policy: RequestPolicy + TerminalDispatcher, - ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, - { - let context = request.lifecycle_context(); - self.run_result(context, request, policy, provider_call) - .await - } } async fn failure( @@ -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; diff --git a/litellm-rust/crates/core/src/lifecycle/ocr.rs b/litellm-rust/crates/core/src/lifecycle/ocr.rs index d3eba181b9f..b054ca7baec 100644 --- a/litellm-rust/crates/core/src/lifecycle/ocr.rs +++ b/litellm-rust/crates/core/src/lifecycle/ocr.rs @@ -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; impl Lifecycle { - pub fn new(admission: &OcrRequest, options: Options) -> Result, Error> { + pub fn new( + admission: &OcrAdmissionRequest, + options: Options, + ) -> Result, Error> { >::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); + } +} diff --git a/litellm-rust/crates/core/src/messages/lifecycle.rs b/litellm-rust/crates/core/src/messages/lifecycle.rs index 63210bf3f24..48feef64f1f 100644 --- a/litellm-rust/crates/core/src/messages/lifecycle.rs +++ b/litellm-rust/crates/core/src/messages/lifecycle.rs @@ -238,9 +238,14 @@ pub async fn messages( context: CallLifecycleContext, ) -> ExecutedCall { 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 } diff --git a/litellm-rust/crates/core/src/ocr/common_utils.rs b/litellm-rust/crates/core/src/ocr/common_utils.rs index f2799244071..5e7cbf79feb 100644 --- a/litellm-rust/crates/core/src/ocr/common_utils.rs +++ b/litellm-rust/crates/core/src/ocr/common_utils.rs @@ -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::*; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 4c8401cd1fc..00d1b7d09e0 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -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 { - 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); diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 99e6b7ea8a2..441e26c6b25 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -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> + 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 for OcrLifecycleHooks { +impl RequestPolicy for OcrRequestPolicy { type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; @@ -239,23 +233,6 @@ impl RequestPolicy 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", - } -} diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs deleted file mode 100644 index 320f09ad478..00000000000 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ /dev/null @@ -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); - } -} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index b73ea3831f6..7e0d6fae987 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -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 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( + services: &S, + request: OcrRouteRequest<'_>, + _options: crate::lifecycle::ocr::Options, + context: CallLifecycleContext, +) -> ExecutedCall { + 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 for PreparedOcrPolicy { + type PreCallFuture<'a> + = std::future::Ready> + where + Self: 'a; + type DuringCallFuture<'a> + = std::future::Ready> + 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, diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 8b4ac9b92f5..5c181d080f4 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -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, ) -> Result<(), Error> { let (provider, config) = request_config(request)?; @@ -81,7 +81,7 @@ fn check_admission_capabilities( Ok(()) } -pub fn prepare(request: OcrRequest) -> Result { +pub fn prepare(request: OcrAdmissionRequest) -> Result { 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 }; diff --git a/litellm-rust/crates/core/src/ocr/runtime_types.rs b/litellm-rust/crates/core/src/ocr/runtime_types.rs index d27f6a92188..a5509fa828c 100644 --- a/litellm-rust/crates/core/src/ocr/runtime_types.rs +++ b/litellm-rust/crates/core/src/ocr/runtime_types.rs @@ -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> for OcrRouteRequest<'a> { + fn from(request: OcrRequest<'a>) -> Self { + Self::Native(request) + } +} + +impl From 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, diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 78f8fab5b3e..63a21a51cb4 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -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, pub api_key: Option, @@ -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, pub model: String, + #[serde(default)] pub document_annotation: Option, + #[serde(default)] pub usage_info: Option, pub object: String, + #[serde(flatten)] pub extra_fields: Map, + #[serde(default)] pub provider_native_response: Option, } diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index 5f29dbb143e..007bec7b710 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -330,6 +330,8 @@ mod tests { ), (), &instrumentation, + &instrumentation, + &crate::lifecycle::SystemClock, |_| async { Ok::<(), Error>(()) }, ) .await; diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index c200db6c951..4a6be72964e 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -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 { + 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 { diff --git a/litellm-rust/crates/core/tests/ocr_lifecycle.rs b/litellm-rust/crates/core/tests/ocr_lifecycle.rs index eb5b167200a..592b59ebf1b 100644 --- a/litellm-rust/crates/core/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/core/tests/ocr_lifecycle.rs @@ -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 { + 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]; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index dedfae4108f..df663885e6e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -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) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 2fa9599c4c9..4db0cc70d8e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -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>, logging: Option>, prepared: Option, + terminal: Option, } #[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> .collect() } -fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult { +fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult { 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, asynchronous: bool, } @@ -399,6 +405,7 @@ fn prepare(py: Python<'_>, arguments: Py, 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) -> PyResult { fn send(py: Python<'_>, state: Py) -> PyResult> { 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) -> PyResult> { 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::()?; + 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::()?; let response = finish(py, fields.unbind()); drop(state); response } +#[pyfunction] +fn terminal_record(py: Python<'_>, state: Py) -> PyResult> { + 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) -> PyResult> { 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::() + .unwrap(), + 10.25 + ); + assert_eq!( + timing + .get_item("end_time") + .unwrap() + .extract::() + .unwrap(), + 12.5 + ); + }); + } + #[test] fn callback_decline_is_terminal_and_identity_is_reused() { Python::initialize(); diff --git a/litellm/rust_bridge/_lifecycle.py b/litellm/rust_bridge/_lifecycle.py index 75aa86a24d4..72cce9d9ffb 100644 --- a/litellm/rust_bridge/_lifecycle.py +++ b/litellm/rust_bridge/_lifecycle.py @@ -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) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 2d9e8aec9f3..fe52cca72ef 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -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__))