diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 5b7180d0148..c5470d6845f 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -51,10 +51,14 @@ function per top-level route, mirroring the core entrypoints. ## Checks -### Native OCR Boundary +### Private Native OCR Proof -`LITELLM_RUST=1` selects native OCR before Python provider preparation. With Rust -disabled, the existing Python execution and authentication paths are unchanged +Public `litellm.ocr` and `litellm.aocr` always use the existing Python lifecycle, +including when `litellm.rust(True)` or `LITELLM_RUST=1` enables other Rust paths. +Native OCR remains a private proof until full lifecycle parity is established. +Only tests requesting the private `native_ocr` fixture replace those public +functions with test-only route selection: Rust enabled calls the native bridge, +and Rust disabled calls the captured production Python functions The bridge retains the complete call argument dictionary as a Python object, including opaque callback and metadata objects. It creates callback-visible @@ -65,7 +69,7 @@ chat messages using the existing Rust transform. Rust performs provider preparat encoding, HTTP and response normalization. Python continues to dispatch existing logging operations and construct the public response object -This is an opt-in implementation scaffold, not full OCR parity. Azure Mistral and +This is a private implementation scaffold, not full OCR parity. Azure Mistral and Vertex Mistral accept inline data URIs with supplied keys/tokens, native environment keys or auth headers. Azure also accepts a supplied `azure_ad_token`. Vertex DeepSeek uses its existing chat request and OCR response transforms. Cloud @@ -74,9 +78,9 @@ HTTP document URL conversion fails only for configs requiring data URIs. Azure Document Intelligence selects its own config but fails at the polling capability check before sending a billable analyze request. Cohere transforms, file inputs, streaming, native response format and compression remain unsupported. -An enabled but missing native extension also fails; -neither case falls back to Python execution. Transport failures currently use a -generic error rather than the SDK's timeout-specific exception +Direct private bridge calls with a missing native extension also fail; +neither case falls back to Python execution within the private route. Transport +failures currently use a generic error rather than the SDK's timeout-specific exception Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust changes. That list is the single source of truth and matches what GitHub Actions @@ -100,7 +104,7 @@ make lint-rust-python-fixtures `lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without syncing the project environment -Run the native OCR acceptance gate from the repository root: +Run the private native OCR proof gate from the repository root: ```bash make test-rust-ocr @@ -126,7 +130,30 @@ invocation context and ownership against Python behavior, including existing LiteLLM components. Short synthetic pre-call contracts use Rust-owned table-driven cases with inline Python callbacks; larger component scenarios share Python fixtures. These generic proofs complement, rather than replace, native OCR -acceptance tests +private proof tests + +The standard-library-only tests in +`crates/python-interop/tests/callback_patterns.rs` define small inline Python +callbacks, with Rust controlling invocation, ownership and assertions. They +compare Python-reference and Rust-retained calls using both direct and awaited +invocation. They model the behavior +groups in the callback use-case inventory: live versus serialized queues, +mutation before an error, ignored returns, identity-based redaction, block-state +stash, background writes after return, parallel live data versus snapshots, and +shallow/deep copies with uncopyable-value fallback. Copy controls deliberately +produce different observations; event gates establish ordering without sleeps. +Existing synthetic lifecycle cases also cover streams, context and cancellation + +Run this matrix without LiteLLM, vendor SDKs, credentials or services: + +```bash +cargo test --manifest-path litellm-rust/Cargo.toml -p litellm-python-interop --test callback_patterns +``` + +These are behavioral models, not tests of vendor authentication, delivery or +production dispatcher policy. The optional component and integration fixtures +exercise existing LiteLLM implementations with fake transports and credentials +as supplementary coverage; run them with `make test-rust-python` The callback lifecycle scenarios use `#[serial(python_interpreter)]` to isolate CPython GC and interpreter-wide diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs new file mode 100644 index 00000000000..3868cd83b50 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -0,0 +1,600 @@ +use crate::Error; + +use super::{OcrRequest, prepare}; + +#[derive(Debug)] +pub enum NativeOutcome { + Completed(T), + Declined(Decline), +} + +#[derive(Debug, PartialEq, Eq)] +pub struct Decline(&'static str); + +impl Decline { + pub fn reason(&self) -> &'static str { + self.0 + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CredentialMethod { + #[default] + Configured, + Acquisition, +} + +#[derive(Default)] +pub struct Options { + pub asynchronous: bool, + pub internal_call: bool, + pub call_id: Option, + pub trace_id: Option, + pub credential_method: CredentialMethod, +} + +#[derive(Debug, PartialEq, Eq)] +pub struct Identity { + pub requested_model: String, + pub call_id: String, + pub trace_id: Option, + pub generated_call_id: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome { + Success, + Failure, + Abort, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Setup, + DeploymentPre, + Prepare, + Send, + DeploymentSuccess, + DeploymentFailure, + SyncSuccess, + AsyncSuccess, + SyncSuccessIfNeeded, + SyncFailure, + AsyncFailure, + Restore, + Complete(Outcome), +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct Observations { + pub logger_available: bool, + pub has_fallbacks: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ErrorDisposition { + Preserve, + Replace, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Transition { + pub operation: Operation, + pub error: ErrorDisposition, +} + +#[derive(Debug)] +pub struct Lifecycle { + operation: Operation, + outcome: Outcome, + asynchronous: bool, + internal_call: bool, + identity: Identity, +} + +impl Lifecycle { + pub fn new(admission: &OcrRequest, options: Options) -> Result, Error> { + match prepare::admission_capabilities(admission) { + Err(Error::Unsupported(reason)) => return Ok(NativeOutcome::Declined(Decline(reason))), + Err(error) => return Err(error), + Ok(()) => {} + } + if options.credential_method == CredentialMethod::Acquisition { + return Ok(NativeOutcome::Declined(Decline( + "OCR credential acquisition", + ))); + } + let generated_call_id = options.call_id.is_none(); + let call_id = options.call_id.unwrap_or_else(|| { + let id = (rand::random::() & !(0xf000_u128 << 64 | 0xc000_u128 << 48)) + | (0x4000_u128 << 64 | 0x8000_u128 << 48); + let hex = format!("{id:032x}"); + format!( + "{}-{}-{}-{}-{}", + &hex[..8], + &hex[8..12], + &hex[12..16], + &hex[16..20], + &hex[20..] + ) + }); + Ok(NativeOutcome::Completed(Self { + operation: Operation::Setup, + outcome: Outcome::Success, + asynchronous: options.asynchronous, + internal_call: options.internal_call, + identity: Identity { + requested_model: admission.model.clone(), + call_id, + trace_id: options.trace_id, + generated_call_id, + }, + })) + } + + pub fn operation(&self) -> Operation { + self.operation + } + + pub fn identity(&self) -> &Identity { + &self.identity + } + + pub fn advance( + &mut self, + outcome: Outcome, + observations: Observations, + ) -> Result { + use Operation::*; + + if matches!(self.operation, Complete(_)) { + return Err(Error::InvalidRequest( + "OCR lifecycle is already complete".into(), + )); + } + let failure = if observations.logger_available && !(self.asynchronous && self.internal_call) + { + SyncFailure + } else { + Restore + }; + let error = if outcome != Outcome::Success && self.operation != DeploymentFailure { + self.outcome = outcome; + ErrorDisposition::Replace + } else { + ErrorDisposition::Preserve + }; + self.operation = match (self.operation, outcome) { + (Restore, _) => Complete(self.outcome), + (DeploymentFailure, _) => failure, + (_, Outcome::Abort) => Restore, + (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, + (Prepare | Send, Outcome::Failure) if self.asynchronous => DeploymentFailure, + (_, Outcome::Failure) => failure, + (Setup, Outcome::Success) if self.asynchronous => DeploymentPre, + (Setup | DeploymentPre, Outcome::Success) => Prepare, + (Prepare, Outcome::Success) => Send, + (Send, Outcome::Success) if self.asynchronous => DeploymentSuccess, + (Send, Outcome::Success) => SyncSuccess, + (DeploymentSuccess, Outcome::Success) => { + if self.internal_call || observations.has_fallbacks { + SyncSuccessIfNeeded + } else { + AsyncSuccess + } + } + (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, + (SyncFailure, Outcome::Success) if self.asynchronous => AsyncFailure, + (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { + Restore + } + (Complete(_), _) => unreachable!(), + }; + Ok(Transition { + operation: self.operation, + error, + }) + } +} + +#[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 420c4d222a0..b73ea3831f6 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,3 +1,4 @@ +pub mod lifecycle; pub mod prepare; pub mod transformation; pub mod types; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index c2d86a3c01d..d6e11c6e543 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -6,13 +6,15 @@ use crate::Error; use crate::providers::azure_ai::ocr::transformation as azure_ai; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; use crate::providers::vertex_ai::ocr::transformation as vertex_ai; -use crate::routing_utils::provider::get_custom_llm_provider; +use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::transformation::{OcrProviderConfig, OcrResponseHandling}; use super::types::OcrRequest; pub use super::types::PreparedOcr; -pub fn prepare(request: OcrRequest) -> Result { +fn request_config( + request: &OcrRequest, +) -> Result<(CustomLlmProvider<'_>, &'static dyn OcrProviderConfig), Error> { match request.request_format.as_deref() { None | Some("litellm") => {} Some("native") => return Err(Error::Unsupported("native OCR request format")), @@ -33,10 +35,58 @@ pub fn prepare(request: OcrRequest) -> Result { .filter(|timeout| !timeout.is_zero()) .ok_or_else(|| Error::InvalidRequest("timeout must be positive and finite".into()))?; + Ok((provider, config)) +} + +pub(super) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> { + check_admission_capabilities(request, &|key| std::env::var(key).ok()) +} + +fn check_admission_capabilities( + request: &OcrRequest, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result<(), Error> { + let (provider, config) = request_config(request)?; + validate_capabilities(config)?; + request + .document + .validate(config.requires_data_uri_document())?; + if request.stream { + return Err(Error::Unsupported("OCR streaming response handling")); + } + if let Some(operation) = config.credential_acquisition_operation() { + let supplied = request + .api_key + .as_deref() + .is_some_and(|key| !key.trim().is_empty()) + || crate::http_utils::has_header(&request.extra_headers, "authorization"); + let configured = match provider.custom_llm_provider { + "azure_ai" => { + crate::http_utils::has_header(&request.extra_headers, "api-key") + || request + .azure_ad_token + .as_deref() + .is_some_and(|key| !key.trim().is_empty()) + || env_lookup("AZURE_AI_API_KEY").is_some_and(|key| !key.trim().is_empty()) + } + "vertex_ai" => ["VERTEX_AI_API_KEY", "VERTEXAI_API_KEY"] + .into_iter() + .any(|name| env_lookup(name).is_some_and(|key| !key.trim().is_empty())), + _ => false, + }; + if !supplied && !configured { + return Err(Error::Unsupported(operation)); + } + } + Ok(()) +} + +pub fn prepare(request: OcrRequest) -> Result { + let (provider, config) = request_config(&request)?; let env_lookup = |key: &str| std::env::var(key).ok(); let headers = config .validate_credentials( - request.extra_headers, + request.extra_headers.clone(), request.api_key.as_deref(), request.azure_ad_token.as_deref(), &env_lookup, @@ -55,11 +105,11 @@ pub fn prepare(request: OcrRequest) -> Result { return Err(Error::Unsupported("OCR streaming response handling")); } let url_params = [ - ("vertex_project", request.vertex_project), - ("vertex_location", request.vertex_location), + ("vertex_project", request.vertex_project.as_ref()), + ("vertex_location", request.vertex_location.as_ref()), ] .into_iter() - .filter_map(|(name, value)| value.map(|value| (name.into(), Value::String(value)))) + .filter_map(|(name, value)| value.map(|value| (name.into(), Value::String(value.clone())))) .collect(); let url = config.complete_url( request.api_base.as_deref(), @@ -72,7 +122,7 @@ pub fn prepare(request: OcrRequest) -> Result { if !matches!(parsed_url.scheme(), "http" | "https") || parsed_url.host_str().is_none() { return Err(Error::InvalidRequest("invalid OCR API URL".into())); } - let document = serde_json::to_value(request.document) + let document = serde_json::to_value(&request.document) .map_err(|_| Error::InvalidRequest("could not project OCR document".into()))?; let template = config.transform_ocr_request(provider.model, document, Map::new())?; if template.files.is_some() { @@ -113,3 +163,133 @@ pub(super) fn provider_config( _ => Err(Error::Unsupported("OCR provider")), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::ocr::types::OcrDocument; + + fn request() -> OcrRequest { + OcrRequest { + model: "mistral/mistral-ocr-latest".into(), + custom_llm_provider: None, + api_key: Some("test-key".into()), + api_base: Some("not a URL".into()), + extra_headers: vec![], + timeout_seconds: 2.0, + request_format: None, + document: OcrDocument::DocumentUrl { + document_url: "https://example.test/document.pdf".into(), + }, + azure_ad_token: None, + vertex_project: None, + vertex_location: None, + stream: false, + } + } + + #[rstest::rstest] + #[case("openai/model", Some("native"), "native OCR request format")] + #[case("openai/model", None, "OCR provider")] + #[case( + "azure_ai/doc-intelligence/prebuilt-read", + None, + "Azure Document Intelligence OCR polling" + )] + #[case( + "vertex_ai/mistral-ocr-latest", + None, + "OCR HTTP document URL to data URI conversion" + )] + #[case("mistral/mistral-ocr-latest", None, "OCR streaming response handling")] + fn admission_preserves_unsupported_error_order( + #[case] model: &str, + #[case] format: Option<&str>, + #[case] expected: &str, + ) { + let request = OcrRequest { + model: model.into(), + request_format: format.map(str::to_owned), + stream: true, + ..request() + }; + assert!( + matches!(check_admission_capabilities(&request, &|_| None), Err(Error::Unsupported(message)) if message == expected) + ); + assert!( + matches!(prepare(request), Err(Error::Unsupported(message)) if message == expected) + ); + } + + #[test] + fn admission_leaves_url_preparation_and_revalidation_until_prepare() { + let request = request(); + assert!(check_admission_capabilities(&request, &|_| None).is_ok()); + assert!(matches!(prepare(request), Err(Error::InvalidRequest(_)))); + + let mut request = self::request(); + assert!(check_admission_capabilities(&request, &|_| None).is_ok()); + request.timeout_seconds = 0.0; + request.stream = true; + assert!(matches!( + check_admission_capabilities(&request, &|_| None), + Err(Error::InvalidRequest(_)) + )); + assert!(matches!(prepare(request), Err(Error::InvalidRequest(_)))); + } + + #[test] + fn admission_does_not_validate_credentials_or_resolve_headers() { + let request = OcrRequest { + api_key: None, + ..request() + }; + assert!( + check_admission_capabilities(&request, &|_| panic!( + "Mistral admission needs no key lookup" + )) + .is_ok() + ); + assert!( + MISTRAL_OCR_CONFIG + .validate_credentials(vec![], None, None, &|_| None) + .is_err() + ); + } + + #[rstest::rstest] + #[case("azure_ai/model", "AZURE_AI_API_KEY")] + #[case("vertex_ai/model", "VERTEX_AI_API_KEY")] + #[case("vertex_ai/model", "VERTEXAI_API_KEY")] + fn admission_checks_credential_method_using_only_inert_configuration( + #[case] model: &str, + #[case] env_name: &str, + ) { + let request = OcrRequest { + model: model.into(), + api_key: None, + document: OcrDocument::ImageUrl { + image_url: "data:image/png;base64,AA==".into(), + }, + ..request() + }; + assert!(matches!( + check_admission_capabilities(&request, &|_| None), + Err(Error::Unsupported(_)) + )); + assert!(matches!( + check_admission_capabilities(&request, &|_| Some(" ".into())), + Err(Error::Unsupported(_)) + )); + assert!( + check_admission_capabilities(&request, &|name| (name == env_name) + .then(|| "configured".into())) + .is_ok() + ); + let request = OcrRequest { + extra_headers: vec![("AUTHORIZATION".into(), "inert".into())], + ..request + }; + assert!(check_admission_capabilities(&request, &|_| None).is_ok()); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 7ff84273fbb..e626b6f3d2f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -3,14 +3,19 @@ //! 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::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 litellm_python_interop::{ + InvocationMode, InvocationOutcome, PreparedCall, Pythonized, from_py, to_py, +}; use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::pyclass::{PyTraverseError, PyVisit}; use pyo3::sync::PyOnceLock; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyTuple}; use serde_json::Value; use crate::errors::core_error_to_pyerr; @@ -50,7 +55,7 @@ impl OcrState { fn ocr_error_to_pyerr(py: Python<'_>, error: Error, model: &str, provider: &str) -> PyErr { let status = match error { - Error::Unsupported(message) => return PyNotImplementedError::new_err(message), + Error::Unsupported(message) => return PyRuntimeError::new_err(message), Error::Auth(_) => 401, Error::Http { status, .. } => status, Error::Network(_) | Error::Connect(_) => { @@ -120,24 +125,11 @@ fn header_pairs(headers: &Bound<'_, PyDict>) -> PyResult> .collect() } -#[pyfunction] -#[pyo3(signature = (arguments, asynchronous=false))] -fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResult> { - let bag = arguments.bind(py); +fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult { let document = bag .get_item("document")? .ok_or_else(|| PyValueError::new_err("OCR requires document"))? .cast_into::()?; - if scalar(&document, "type")?.as_deref() == Some("file") { - return Err(ocr_error_to_pyerr( - py, - Error::Unsupported( - "Native OCR does not support file documents; pass a document_url or image_url dict", - ), - "", - "", - )); - } let timeout_seconds = match bag.get_item("timeout")?.filter(|value| !value.is_none()) { None => py .import("litellm.constants")? @@ -163,9 +155,9 @@ fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResul document_input.set_item(name, value)?; } } - let prepared = litellm_core::ocr::prepare::prepare(OcrRequest { - model: model.clone(), - custom_llm_provider: custom_llm_provider.clone(), + Ok(OcrRequest { + model, + custom_llm_provider, api_key: scalar(bag, "api_key")?, api_base: scalar(bag, "api_base")?, extra_headers, @@ -182,20 +174,179 @@ fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResul .transpose()? .unwrap_or(false), }) - .map_err(|error| { - let resolved = get_custom_llm_provider(&model, custom_llm_provider.as_deref()); - ocr_error_to_pyerr( - py, - error, - resolved - .as_ref() - .map_or(model.as_str(), |value| value.model), - resolved - .as_ref() - .map_or("", |value| value.custom_llm_provider), - ) - })?; +} +fn request_error_to_pyerr( + py: Python<'_>, + error: Error, + model: &str, + custom_llm_provider: Option<&str>, +) -> PyErr { + let resolved = get_custom_llm_provider(model, custom_llm_provider); + ocr_error_to_pyerr( + py, + error, + resolved.as_ref().map_or(model, |value| value.model), + resolved + .as_ref() + .map_or("", |value| value.custom_llm_provider), + ) +} + +#[pyclass] +struct OcrLifecycle { + machine: Lifecycle, + asynchronous: bool, +} + +#[pymethods] +impl OcrLifecycle { + #[new] + fn new( + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + asynchronous: bool, + internal_call: bool, + ) -> PyResult { + let request = decode_request(py, arguments)?; + let logger = arguments + .get_item("litellm_logging_obj")? + .filter(|value| !value.is_none()); + let identity = |name: &str| -> PyResult> { + if let Some(logger) = &logger { + match logger.getattr(name) { + Ok(value) => { + if let Ok(value) = value.extract::() { + return Ok(Some(value)); + } + } + Err(error) + if !error.is_instance_of::(py) => + { + return Err(error); + } + _ => {} + } + } + scalar(arguments, name) + }; + let machine = Lifecycle::new( + &request, + Options { + asynchronous, + internal_call, + call_id: identity("litellm_call_id")?, + trace_id: identity("litellm_trace_id")?, + ..Options::default() + }, + ) + .map_err(|error| { + request_error_to_pyerr( + py, + error, + &request.model, + request.custom_llm_provider.as_deref(), + ) + })?; + match machine { + NativeOutcome::Completed(machine) => Ok(Self { + machine, + asynchronous, + }), + NativeOutcome::Declined(decline) => { + Err(PyNotImplementedError::new_err(decline.reason())) + } + } + } + + fn identity(&self) -> (String, Option) { + let identity = self.machine.identity(); + (identity.call_id.clone(), identity.trace_id.clone()) + } + + fn advance( + &mut self, + outcome: u8, + logger_available: bool, + has_fallbacks: bool, + ) -> PyResult { + let outcome = match outcome { + 0 => Outcome::Success, + 1 => Outcome::Failure, + _ => Outcome::Abort, + }; + self.machine + .advance( + outcome, + Observations { + logger_available, + has_fallbacks, + }, + ) + .map(|transition| transition.error == ErrorDisposition::Replace) + .map_err(core_error_to_pyerr) + } + + fn complete(&self) -> Option { + match self.machine.operation() { + Operation::Complete(outcome) => Some(outcome == Outcome::Success), + _ => None, + } + } +} + +#[pyfunction] +fn invoke( + py: Python<'_>, + machine: Py, + host: Py, +) -> PyResult<(bool, Py)> { + let (operation, asynchronous) = { + let machine = machine.borrow(py); + (machine.machine.operation(), machine.asynchronous) + }; + let (method, mode) = match operation { + Operation::Setup => ("setup", InvocationMode::Direct), + Operation::DeploymentPre => ("deployment_pre", InvocationMode::Await), + Operation::Prepare => ("prepare", InvocationMode::Direct), + Operation::Send if asynchronous => ("send", InvocationMode::Await), + Operation::Send => ("send_sync", InvocationMode::Direct), + Operation::DeploymentSuccess => ("deployment_success", InvocationMode::Await), + Operation::DeploymentFailure => ("deployment_failure", InvocationMode::Await), + Operation::SyncSuccess => ("sync_success", InvocationMode::Direct), + Operation::AsyncSuccess => ("async_success", InvocationMode::Direct), + Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", InvocationMode::Direct), + Operation::SyncFailure => ("sync_failure", InvocationMode::Direct), + Operation::AsyncFailure => ("async_failure", InvocationMode::Await), + Operation::Restore => ("restore", InvocationMode::Direct), + Operation::Complete(_) => return Err(PyRuntimeError::new_err("OCR lifecycle is complete")), + }; + let call = PreparedCall::new( + mode, + host.getattr(py, method)?, + PyTuple::empty(py).unbind(), + None, + ); + match call.invoke(py)? { + InvocationOutcome::Returned(value) => Ok((false, value)), + InvocationOutcome::Awaitable(value) => Ok((true, value)), + } +} + +#[pyfunction] +#[pyo3(signature = (arguments, asynchronous=false))] +fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResult> { + let bag = arguments.bind(py); + let request = decode_request(py, bag)?; + let model = request.model.clone(); + let custom_llm_provider = request.custom_llm_provider.clone(); + let prepared = litellm_core::ocr::prepare::prepare(request).map_err(|error| { + request_error_to_pyerr(py, error, &model, custom_llm_provider.as_deref()) + })?; + let document = bag + .get_item("document")? + .ok_or_else(|| PyValueError::new_err("OCR requires document"))? + .cast_into::()?; let body = to_py(py, &prepared.body)? .into_bound(py) .cast_into::()?; @@ -355,64 +506,128 @@ fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { let module = PyModule::from_code( py, c"from datetime import datetime -from litellm.rust_bridge.ocr import invoke_terminal +from litellm import utils +from litellm.types.utils import CallTypes +from litellm.rust_bridge.ocr import initialize_logging, invoke_terminal + +class Host: + def __init__(self, arguments, asynchronous): + self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get()) + self.arguments = arguments + self.current = arguments + self.asynchronous = asynchronous + self.logger = arguments.get('litellm_logging_obj') + self.state = None + self.response = None + self.error = None + self.start = datetime.now() + self.end = None + + def setup(self): + call_id, trace_id = self.machine.identity() + self.arguments['litellm_call_id'] = call_id + self.arguments['litellm_trace_id'] = trace_id + self.logger = initialize_logging(self.arguments, self.asynchronous) + self.arguments['litellm_logging_obj'] = self.logger + + async def deployment_pre(self): + modified = await utils.async_pre_call_deployment_hook(self.current, 'aocr') + if modified is not None: + self.current = modified + self.current['litellm_logging_obj'] = self.logger + call_id, trace_id = self.machine.identity() + self.current['litellm_call_id'] = call_id + self.current['litellm_trace_id'] = trace_id + + def prepare(self): + self.state = _prepare(self.current, self.asynchronous) + + def send_sync(self): + self.response = _send_sync(self.state) + self.end = datetime.now() + + async def send(self): + self.response = _finish(await _send(self.state)) + self.end = datetime.now() + + async def deployment_success(self): + self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aocr) + + async def deployment_failure(self): + 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) + + def sync_success(self): + return self.terminal('sync_success', self.response) + + def async_success(self): + return self.terminal('async_success', self.response) + + def sync_success_if_needed(self): + return self.terminal('sync_success_if_needed', self.response) + + def sync_failure(self): + return self.terminal('sync_failure', self.error) + + def async_failure(self): + return self.terminal('async_failure', self.error) + + def restore(self): + utils._restore_correlation_context_if_supported(self.logger) + + def advance(self, outcome, error=None): + if error is not None and self.end is None: + self.end = datetime.now() + if self.logger is None: + self.logger = self.arguments.get('litellm_logging_obj') + replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None) + if replace: + self.error = error + + def result(self): + if self.machine.complete(): + return self.response + raise self.error def drive_sync(arguments): - start = datetime.now() - state = None - try: - state = _prepare(arguments, False) - response = _send_sync(state) - except Exception as error: - logger = arguments.get('litellm_logging_obj') - if state is not None or (logger is not None and not isinstance(error, NotImplementedError)): - end = datetime.now() - for action in _sync_failure: - invoke_terminal(action, (arguments, state), logger, error, start, end) - raise - end = datetime.now() - for action in _sync_success: - invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end) - return response + host = Host(arguments, False) + while host.machine.complete() is None: + try: + _invoke(host.machine, host) + except Exception as error: + host.advance(1, error) + except BaseException as error: + host.advance(2, error) + else: + host.advance(0) + return host.result() async def drive(arguments): - start = datetime.now() - state = None - try: - state = _prepare(arguments, True) - response = _finish(await _send(state)) - except Exception as error: - logger = arguments.get('litellm_logging_obj') - if state is not None or (logger is not None and not isinstance(error, NotImplementedError)): - end = datetime.now() - for action in _async_failure: - pending = invoke_terminal(action, (arguments, state), logger, error, start, end) - if pending is not None: - await pending - raise - end = datetime.now() - for action in _async_success: - invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end) - return response + host = Host(arguments, True) + while host.machine.complete() is None: + try: + awaiting, value = _invoke(host.machine, host) + if awaiting: + await value + except Exception as error: + host.advance(1, error) + except BaseException as error: + host.advance(2, error) + else: + host.advance(0) + return host.result() ", c"ocr_driver.py", c"_ocr_driver", )?; + module.add("_Lifecycle", py.get_type::())?; + module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?; module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; module.add("_send", wrap_pyfunction!(send, &module)?)?; module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; module.add("_finish", wrap_pyfunction!(finish, &module)?)?; - for (name, asynchronous, success) in [ - ("_sync_success", false, true), - ("_async_success", true, true), - ("_sync_failure", false, false), - ("_async_failure", true, false), - ] { - module.add( - name, - litellm_core::ocr::terminal_callbacks(asynchronous, success).to_vec(), - )?; - } Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) } @@ -519,6 +734,88 @@ mod tests { }); } + #[test] + fn callback_decline_is_terminal_and_identity_is_reused() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "ocr_test").unwrap(); + module + .add_function(wrap_pyfunction!(ocr, &module).unwrap()) + .unwrap(); + module + .add_function(wrap_pyfunction!(aocr, &module).unwrap()) + .unwrap(); + let globals = PyDict::new(py); + globals.set_item("native", module).unwrap(); + py.run( + c" +import asyncio +import contextvars +import threading +from datetime import datetime + +marker = contextvars.ContextVar('terminal_marker') + +class Logger: + litellm_call_id = 'supplied-call' + litellm_trace_id = 'supplied-trace' + + def update_from_kwargs(self, **values): + assert values['kwargs']['litellm_call_id'] == self.litellm_call_id + assert values['kwargs']['litellm_trace_id'] == self.litellm_trace_id + assert threading.get_ident() == self.thread + marker.set('update') + raise self.original + + def failure_handler(self, error, trace, start, end): + assert error is self.original + assert marker.get() == 'update' + assert start <= end <= datetime.now() + self.end = end + self.calls.append('failure') + + async def async_failure_handler(self, error, trace, start, end): + await asyncio.sleep(0) + assert asyncio.current_task() is self.task + assert marker.get() == 'update' + assert error is self.original + assert end is self.end + self.calls.append('async_failure') + + def _restore_correlation_context(self): + self.calls.append('restore') + +async def exercise(): + for asynchronous in (False, True): + logger = Logger() + logger.thread = threading.get_ident() + logger.task = asyncio.current_task() + logger.calls = [] + logger.original = NotImplementedError('callback declined, not admission') + arguments = dict(model='mistral/mistral-ocr-latest', api_key='test-key', timeout=1.0, + document={'type': 'document_url', 'document_url': 'https://example.test/doc.pdf'}, + litellm_logging_obj=logger) + try: + if asynchronous: + await native.aocr(arguments) + else: + native.ocr(arguments) + except NotImplementedError as error: + assert error is logger.original + else: + raise AssertionError('callback exception was lost') + assert logger.calls == (['failure', 'async_failure', 'restore'] if asynchronous else ['failure', 'restore']) + assert arguments['litellm_call_id'] == 'supplied-call' + assert arguments['litellm_trace_id'] == 'supplied-trace' + +asyncio.run(exercise()) +", + Some(&globals), + Some(&globals), + ).unwrap(); + }); + } + #[test] fn native_send_owns_state_without_the_python_driver() { Python::initialize(); @@ -733,7 +1030,8 @@ class Logger: assert asyncio.current_task() is caller assert threading.get_ident() == caller_thread assert marker.get() == 'caller' - assert values['kwargs'] is arguments + assert values['kwargs'] is not arguments + assert values['kwargs']['opaque'] is arguments['opaque'] self.calls.append('update') marker.set('updated') diff --git a/litellm-rust/crates/python-interop/tests/callback_controls.rs b/litellm-rust/crates/python-interop/tests/callback_controls.rs new file mode 100644 index 00000000000..7261c84ecf3 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/callback_controls.rs @@ -0,0 +1,227 @@ +use std::ffi::CStr; + +use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use rstest::rstest; +use serial_test::serial; + +#[path = "support/mod.rs"] +mod support; + +use support::python::{InitializedPython, initialized_python, item, scope}; + +enum ControlCall { + Reference(Py), + Prepared(PreparedCall), +} + +impl ControlCall { + fn new( + py: Python<'_>, + callback: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Option>, + retained: bool, + mode: InvocationMode, + ) -> PyResult { + if retained { + return Ok(Self::Prepared(PreparedCall::new( + mode, + callback.unbind(), + args.unbind(), + kwargs.map(Bound::unbind), + ))); + } + let factory = py + .import("callback_lifecycle")? + .getattr("ReferenceFactory")? + .call0()?; + Ok(Self::Reference( + factory + .call_method1( + "prepare", + (callback, args, kwargs, mode == InvocationMode::Await), + )? + .unbind(), + )) + } + + fn invoke(&self, py: Python<'_>, mode: InvocationMode) -> PyResult> { + match self { + Self::Reference(owner) => owner.call_method0(py, "invoke"), + Self::Prepared(call) => match call.invoke(py)? { + InvocationOutcome::Returned(value) => { + assert_eq!(mode, InvocationMode::Direct); + Ok(value) + } + InvocationOutcome::Awaitable(value) => { + assert_eq!(mode, InvocationMode::Await); + Ok(value) + } + }, + } + } +} + +#[rstest] +#[case::original_graph(c"(positional, keywords)", (true, true, true), (1, 1, 1))] +#[case::rebuilt_envelopes(c"(tuple(value for value in positional), dict(keywords))", (true, true, true), (1, 1, 1))] +#[case::shallow_payload(c"((copy.copy(positional[0]),), keywords)", (false, true, true), (0, 1, 1))] +#[case::copied_graph_keeps_cross_argument_alias(c"copy.deepcopy((positional, keywords))", (false, false, true), (0, 0, 0))] +#[case::separate_copies_break_cross_argument_alias(c"(copy.deepcopy(positional), copy.deepcopy(keywords))", (false, false, false), (0, 0, 0))] +#[serial(python_interpreter)] +fn argument_copy_boundaries_determine_identity_and_mutation_visibility( + initialized_python: &InitializedPython, + #[case] transform: &CStr, + #[case] expected_identity: (bool, bool, bool), + #[case] live_stages: (u8, u8, u8), + #[values(false, true)] retained: bool, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + initialized_python.attach(|py| { + let globals = scope( + py, + c" +import copy +from types import SimpleNamespace + +nested = SimpleNamespace(stage=0) +original = SimpleNamespace(nested=nested, stage=0) +positional, keywords = (original,), {'alias': nested} + +def observe(value, *, alias): + return value, alias, (value.stage, value.nested.stage, alias.stage) + +async def observe_async(value, *, alias): + return observe(value, alias=alias) +", + )?; + let transformed = py.eval(transform, Some(&globals), None)?; + let call = ControlCall::new( + py, + item( + &globals, + if mode == InvocationMode::Await { + "observe_async" + } else { + "observe" + }, + ), + transformed.get_item(0)?.cast_into::()?, + Some(transformed.get_item(1)?.cast_into::()?), + retained, + mode, + )?; + let original = item(&globals, "original"); + let nested = item(&globals, "nested"); + original.setattr("stage", 1)?; + nested.setattr("stage", 1)?; + let pending = call.invoke(py, mode)?; + original.setattr("stage", 2)?; + nested.setattr("stage", 2)?; + let result = if mode == InvocationMode::Await { + let lifecycle = py.import("callback_lifecycle")?; + lifecycle.call_method1( + "run_checked", + (lifecycle.call_method0("ReferenceFactory")?, pending), + )? + } else { + pending.into_bound(py) + }; + drop(call); + let value = result.get_item(0)?; + let alias = result.get_item(1)?; + assert_eq!( + ( + value.is(&original), + value.getattr("nested")?.is(&nested), + value.getattr("nested")?.is(&alias) + ), + expected_identity + ); + let stage = if mode == InvocationMode::Await { 2 } else { 1 }; + assert_eq!( + result.get_item(2)?.extract::<(u8, u8, u8)>()?, + ( + live_stages.0 * stage, + live_stages.1 * stage, + live_stages.2 * stage + ) + ); + Ok(()) + }) +} + +#[rstest] +#[case::original_result(None, (true, true))] +#[case::passthrough_result(Some(c"lambda value: value"), (true, true))] +#[case::shallow_result(Some(c"copy.copy"), (false, true))] +#[case::deep_result(Some(c"copy.deepcopy"), (false, false))] +#[serial(python_interpreter)] +fn result_copy_boundaries_determine_root_and_nested_identity( + initialized_python: &InitializedPython, + #[case] transform: Option<&CStr>, + #[case] expected_identity: (bool, bool), + #[values(false, true)] retained: bool, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + initialized_python.attach(|py| { + let globals = scope( + py, + c" +import copy +from types import SimpleNamespace + +original = SimpleNamespace(nested=SimpleNamespace(stage=0), stage=0) + +def callback(): + return original + +async def callback_async(): + return original +", + )?; + let call = ControlCall::new( + py, + item( + &globals, + if mode == InvocationMode::Await { + "callback_async" + } else { + "callback" + }, + ), + PyTuple::empty(py), + None, + retained, + mode, + )?; + let pending = call.invoke(py, mode)?; + let settled = if mode == InvocationMode::Await { + let lifecycle = py.import("callback_lifecycle")?; + lifecycle.call_method1( + "run_checked", + (lifecycle.call_method0("ReferenceFactory")?, pending), + )? + } else { + pending.into_bound(py) + }; + let result = if let Some(transform) = transform { + py.eval(transform, Some(&globals), None)? + .call1((settled,))? + } else { + settled + }; + drop(call); + let original = item(&globals, "original"); + assert_eq!( + ( + result.is(&original), + result.getattr("nested")?.is(original.getattr("nested")?) + ), + expected_identity + ); + Ok(()) + }) +} diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index 5f50a68a2a2..36cf22c538b 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs @@ -1,8 +1,9 @@ use std::process::Command; use std::time::{Duration, Instant}; +use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyTuple}; use rstest::{fixture, rstest}; use serial_test::{parallel, serial}; @@ -12,7 +13,7 @@ mod callback_owner; #[path = "support/mod.rs"] mod support; -use support::python::{InitializedPython, initialized_python, run_fixture}; +use support::python::{InitializedPython, initialized_python, item, run_fixture}; #[test] fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> { @@ -100,8 +101,6 @@ fn scenario_scope(initialized_python: &InitializedPython) -> Py { #[case::stream_lifecycle("stream_lifecycle")] #[case::sync_stream_lifecycle("sync_stream_lifecycle")] #[case::repeated_ownership("repeated_ownership")] -#[case::retained_field_replacement("retained_field_replacement")] -#[case::queued_graph_ownership("queued_graph_ownership")] #[case::detached_work_after_error("detached_work_after_error")] #[serial(python_interpreter)] fn lifecycle_contract( @@ -176,8 +175,8 @@ fn component_contract( #[case::real_logging_queue_copy_control("real_logging_queue_copy_control")] #[case::real_crowdstrike_translator_identity("real_crowdstrike_translator_identity")] #[case::real_rubrik_block_lifecycle("real_rubrik_block_lifecycle")] -#[case::real_parallel_guardrail_snapshots("real_parallel_guardrail_snapshots")] -#[case::real_purview_sync_background("real_purview_sync_background")] +#[case::real_parallel_guardrail_sharing_and_exception_order("real_parallel_guardrail_snapshots")] +#[case::real_purview_sync_background_and_active_loop("real_purview_sync_background")] #[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] #[serial(python_interpreter)] fn integration_contract( @@ -220,20 +219,6 @@ fn run_scenario_fixture( } #[rstest] -#[case::original_arguments("argument_identity", "identity")] -#[case::envelope_arguments("argument_identity", "envelope")] -#[case::shallow_arguments("argument_identity", "shallow_payload")] -#[case::copied_graph("argument_identity", "deep_graph")] -#[case::independent_copies("argument_identity", "deep_separate")] -#[case::original_read_timing("mutation_timing", "identity")] -#[case::envelope_read_timing("mutation_timing", "envelope")] -#[case::shallow_read_timing("mutation_timing", "shallow_payload")] -#[case::deep_read_timing("mutation_timing", "deep_graph")] -#[case::independent_read_timing("mutation_timing", "deep_separate")] -#[case::original_result("result_identity", "identity")] -#[case::passthrough_result("result_identity", "result_passthrough")] -#[case::shallow_result("result_identity", "result_shallow")] -#[case::deep_result("result_identity", "result_deep")] #[case::retained_lifetime("deferred_lifetime", "identity")] #[case::prepared_ownership("deferred_lifetime", "missing_handoff")] #[case::externally_owned_retained("borrowed_lifetime", "identity")] @@ -305,6 +290,153 @@ fn run_control_fixture( }) } +fn invoke_direct_callback( + py: Python<'_>, + globals: &Bound<'_, PyDict>, + callback: &str, + argument: &str, + retained: bool, +) -> PyResult> { + let args = PyTuple::new(py, [item(globals, argument)])?; + if !retained { + let factory = item(globals, "ReferenceFactory").call0()?; + let owner = factory.call_method1("prepare", (item(globals, callback), args))?; + let result = owner.call_method0("invoke"); + owner.call_method0("close")?; + assert_eq!(factory.getattr("live")?.extract::()?, 0); + return result.map(Bound::unbind); + } + let call = PreparedCall::new( + InvocationMode::Direct, + item(globals, callback).unbind(), + args.unbind(), + None, + ); + match call.invoke(py)? { + InvocationOutcome::Returned(value) => Ok(value), + InvocationOutcome::Awaitable(_) => panic!("direct callback produced an awaitable outcome"), + } +} + +#[rstest] +#[serial(python_interpreter)] +fn retained_field_survives_replacement_and_observes_original_mutations( + scenario_scope: Py, + #[values(false, true)] retained: bool, +) -> PyResult<()> { + Python::attach(|py| { + let globals = scenario_scope.bind(py); + py.run( + c" +original = {'messages': [{'content': 'original'}]} +replacement = {'messages': [{'content': 'replacement'}]} +event = {'payload': original, 'alias': original} +saved = [] + +def retain(value): + saved.append(value['payload']) + +def replace(value): + value['payload'] = replacement + value['alias']['messages'][0]['content'] = 'mutated original' +", + Some(globals), + None, + )?; + for callback in ["retain", "replace"] { + assert!(invoke_direct_callback(py, globals, callback, "event", retained)?.is_none(py)); + } + let original = item(globals, "original"); + let replacement = item(globals, "replacement"); + let event = item(globals, "event"); + let saved = item(globals, "saved").get_item(0)?; + assert!(saved.is(&original)); + assert!(event.get_item("alias")?.is(&original)); + assert!(event.get_item("payload")?.is(&replacement)); + assert_eq!( + saved + .get_item("messages")? + .get_item(0)? + .get_item("content")? + .extract::()?, + "mutated original" + ); + replacement + .get_item("messages")? + .get_item(0)? + .set_item("content", "mutated replacement")?; + assert_eq!( + event + .get_item("payload")? + .get_item("messages")? + .get_item(0)? + .get_item("content")? + .extract::()?, + "mutated replacement" + ); + assert_eq!( + original + .get_item("messages")? + .get_item(0)? + .get_item("content")? + .extract::()?, + "mutated original" + ); + Ok(()) + }) +} + +#[rstest] +#[serial(python_interpreter)] +fn queued_graph_outlives_invocation_and_stays_live_until_serialized( + scenario_scope: Py, + #[values(false, true)] retained: bool, +) -> PyResult<()> { + Python::attach(|py| { + let globals = scenario_scope.bind(py); + py.run( + c" +queue = asyncio.Queue() +sentinel = Value() +reference = weakref.ref(sentinel) +payload = {'sentinel': sentinel, 'nested': {'status': 'queued'}} +snapshot = json.dumps(payload['nested']) + +def enqueue(value): + queue.put_nowait(value) +", + Some(globals), + None, + )?; + assert!(invoke_direct_callback(py, globals, "enqueue", "payload", retained)?.is_none(py)); + py.run(c"del sentinel, payload\ngc.collect()", Some(globals), None)?; + let reference = item(globals, "reference"); + assert!(!reference.call0()?.is_none()); + let queue = item(globals, "queue"); + let queued = queue.call_method0("get_nowait")?; + queued + .get_item("nested")? + .set_item("status", "changed before flush")?; + let json = item(globals, "json"); + let flushed = json.call_method1( + "loads", + (json.call_method1("dumps", (queued.get_item("nested")?,))?,), + )?; + assert_eq!( + flushed.get_item("status")?.extract::()?, + "changed before flush" + ); + let snapshot = json.call_method1("loads", (item(globals, "snapshot"),))?; + assert_eq!(snapshot.get_item("status")?.extract::()?, "queued"); + assert!(queued.get_item("sentinel")?.is(reference.call0()?)); + queue.call_method0("task_done")?; + drop(queued); + py.run(c"gc.collect()", Some(globals), None)?; + assert!(reference.call0()?.is_none()); + Ok(()) + }) +} + #[rstest] #[parallel(python_interpreter)] fn detached_release(initialized_python: &InitializedPython) -> PyResult<()> { diff --git a/litellm-rust/crates/python-interop/tests/callback_patterns.rs b/litellm-rust/crates/python-interop/tests/callback_patterns.rs new file mode 100644 index 00000000000..58e7d5535cc --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/callback_patterns.rs @@ -0,0 +1,720 @@ +use std::ffi::CStr; + +use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use rstest::rstest; +use serial_test::serial; + +#[path = "support/mod.rs"] +mod support; + +use support::python::{InitializedPython, initialized_python, item, run_fixture, scope}; + +#[derive(Clone, Copy, Debug)] +enum Backend { + Python, + PreparedCall, +} + +#[pyfunction] +fn invoke_prepared<'py>( + py: Python<'py>, + callback: Py, + args: Py, + awaited: bool, +) -> PyResult> { + let mode = if awaited { + InvocationMode::Await + } else { + InvocationMode::Direct + }; + let call = PreparedCall::new(mode, callback, args, None); + let outcome = call.invoke(py); + drop(call); + match outcome? { + InvocationOutcome::Returned(value) => { + assert_eq!(mode, InvocationMode::Direct); + Ok(value.into_bound(py)) + } + InvocationOutcome::Awaitable(value) => { + assert_eq!(mode, InvocationMode::Await); + Ok(value.into_bound(py)) + } + } +} + +fn pattern_scope<'py>( + py: Python<'py>, + backend: Backend, + mode: InvocationMode, +) -> PyResult> { + let globals = scope( + py, + c" +import asyncio +import copy +import gc +import json +import threading +import weakref +from unittest import TestCase +from callback_lifecycle import ReferenceFactory, run_checked + +class Value: + pass + +def invoke(callback, *args): + if retained: + return invoke_prepared(callback, args, awaited) + if not awaited: + return callback(*args) + + async def await_callback(): + return await callback(*args) + + return await_callback() + +async def async_call(callback, *args): + async def async_callback(*values): + return callback(*values) + + result = invoke(async_callback if awaited else callback, *args) + return await result if awaited else result + +def call(callback, *args): + if not awaited: + return invoke(callback, *args) + return run_async(async_call(callback, *args)) + +def run_async(scenario): + return run_checked(ReferenceFactory(), scenario) +", + )?; + globals.set_item("invoke_prepared", wrap_pyfunction!(invoke_prepared, py)?)?; + globals.set_item("retained", matches!(backend, Backend::PreparedCall))?; + globals.set_item("awaited", mode == InvocationMode::Await)?; + Ok(globals) +} + +#[track_caller] +fn run_pattern( + python: &InitializedPython, + backend: Backend, + mode: InvocationMode, + parameters: impl for<'py> FnOnce(&Bound<'py, PyDict>) -> PyResult<()>, + scenario: &CStr, +) -> PyResult<()> { + let location = std::panic::Location::caller(); + python.attach(|py| { + let globals = pattern_scope(py, backend, mode)?; + parameters(&globals)?; + run_fixture( + py, + &globals, + scenario.to_str().unwrap(), + &format!("{}:{}", location.file(), location.line()), + )?; + let result = py.run( + c" +if asyncio.iscoroutinefunction(test_pattern): + run_async(test_pattern()) +else: + test_pattern() +", + Some(&globals), + None, + ); + if let Err(error) = &result { + error.print(py); + } + result + }) +} + +#[rstest] +#[case::flush_after_edit_sees_mutation(false, false, false)] +#[case::flush_before_edit_keeps_original(true, false, true)] +#[case::copied_editor_leaves_queue_unchanged(false, true, true)] +#[serial(python_interpreter)] +fn reference_queue_observes_edits_until_serialization_and_owns_arguments( + initialized_python: &InitializedPython, + #[case] flush_before_edit: bool, + #[case] copy_before_edit: bool, + #[case] flushed_has_tools: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| { + globals.set_item("flush_before_edit", flush_before_edit)?; + globals.set_item("copy_before_edit", copy_before_edit)?; + globals.set_item("flushed_has_tools", flushed_has_tools) + }, + c" +def test_pattern(): + queue, snapshots, observations = [], [], [] + payload = {'model_parameters': {'tools': ['lookup'], 'stream': True}} + response = Value() + response_ref = weakref.ref(response) + event = {'payload': payload} + + def enqueue_live(kwargs, result): + queue.append((kwargs['payload'], kwargs, result)) + + def enqueue_snapshot(kwargs, result): + snapshots.append(json.dumps(kwargs['payload'])) + + def pop_tools(kwargs, result): + kwargs['payload']['model_parameters'].pop('tools') + + def observe(kwargs, result): + observations.append('tools' in kwargs['payload']['model_parameters']) + + call(enqueue_live, event, response) + call(enqueue_snapshot, event, response) + early_flush = json.dumps(queue[0][0]) if flush_before_edit else None + edited = copy.deepcopy(event) if copy_before_edit else event + call(pop_tools, edited, response) + call(observe, event, response) + + assert queue[0][0] is payload + assert queue[0][1] is event + assert queue[0][2] is response + assert observations == [copy_before_edit] + assert json.loads(snapshots[0]) == { + 'model_parameters': {'tools': ['lookup'], 'stream': True} + } + + del response, event, payload, edited + gc.collect() + assert response_ref() is not None, 'queued response must outlive invocation' + flushed = json.loads(early_flush if early_flush is not None else json.dumps(queue[0][0])) + assert ('tools' in flushed['model_parameters']) == flushed_has_tools + assert flushed['model_parameters']['stream'] is True + queue.clear() + gc.collect() + assert response_ref() is None, 'clearing the queue must release its response' +", + ) +} + +#[rstest] +#[serial(python_interpreter)] +fn shallow_queue_keeps_nested_aliases_but_not_replaced_fields( + initialized_python: &InitializedPython, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + initialized_python.attach(|py| { + let globals = pattern_scope(py, backend, mode)?; + py.run( + c" +queue = [] +messages = [{'content': 'original'}] +response = Value() +response.messages = messages +event = {'messages': messages, 'status': 'queued'} +replacement = [{'content': 'replacement'}] + +def enqueue(data, result): + queue.append({**data}) + queue.append(result.__dict__) +", + Some(&globals), + None, + )?; + let event = item(&globals, "event"); + let response = item(&globals, "response"); + let messages = item(&globals, "messages"); + item(&globals, "call").call1((item(&globals, "enqueue"), &event, &response))?; + + messages.get_item(0)?.set_item("content", "edited")?; + event.set_item("messages", item(&globals, "replacement"))?; + event.set_item("status", "sent")?; + + let queue = item(&globals, "queue"); + let shallow_entry = queue.get_item(0)?; + let attributes_entry = queue.get_item(1)?; + assert!(!shallow_entry.is(&event)); + assert!(attributes_entry.is(response.getattr("__dict__")?)); + assert!(shallow_entry.get_item("messages")?.is(&messages)); + assert!(attributes_entry.get_item("messages")?.is(&messages)); + let serialized = item(&globals, "json").call_method1("dumps", (&shallow_entry,))?; + assert_eq!( + serde_json::from_str::(serialized.extract::<&str>()?).unwrap(), + serde_json::json!({"messages": [{"content": "edited"}], "status": "queued"}) + ); + assert!( + event + .get_item("messages")? + .is(item(&globals, "replacement")) + ); + Ok(()) + }) +} + +#[rstest] +#[case::ignored_return(false)] +#[case::caught_error_does_not_roll_back(true)] +#[serial(python_interpreter)] +fn sequential_loggers_observe_prior_mutations_and_opaque_state( + initialized_python: &InitializedPython, + #[case] raises: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| globals.set_item("raises", raises), + c" +def test_pattern(): + saved, seen, errors, order = [], [], [], [] + state = {'optional_params': {'tools': ['lookup']}, 'cache_hit': None} + lock = threading.Lock() + failure = RuntimeError('after mutation') + + def retain(kwargs): + order.append('retain') + saved.append(kwargs) + + def mutate(kwargs): + order.append('mutate') + kwargs['optional_params'].pop('tools') + kwargs['cache_hit'] = False + kwargs['flush_lock'] = lock + if raises: + raise failure + return {'replacement': True} + + def observe(kwargs): + order.append('observe') + seen.append((kwargs, kwargs['cache_hit'], kwargs['flush_lock'])) + + for callback in (retain, mutate, observe): + try: + call(callback, state) + except RuntimeError as error: + errors.append(error) + + assert order == ['retain', 'mutate', 'observe'] + assert errors == ([failure] if raises else []) + if raises: + assert errors[0] is failure + assert saved[0] is seen[0][0] is state + assert seen[0][1] is False + assert seen[0][2] is lock + assert state['optional_params'] == {} + assert 'replacement' not in state + with TestCase().assertRaises(TypeError): + json.dumps(state) +", + ) +} + +#[rstest] +#[case::replacement_list(false, false)] +#[case::in_place_redaction(true, false)] +#[case::equal_but_independently_copied_subset_does_not_match(false, true)] +#[serial(python_interpreter)] +fn redaction_matches_message_identity_and_preserves_unscanned_messages( + initialized_python: &InitializedPython, + #[case] redact_in_place: bool, + #[case] copy_subset: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| { + globals.set_item("redact_in_place", redact_in_place)?; + globals.set_item("copy_subset", copy_subset) + }, + c" +def test_pattern(): + first, second = {'content': 'private'}, {'content': 'keep'} + messages = [first, second] + subset = [first] + state = {'messages': messages} + + def redact(full, selected): + if not {id(message) for message in selected} <= {id(message) for message in full}: + return None + if redact_in_place: + selected[0]['content'] = 'masked' + return full + replacements = {id(selected[0]): {'content': 'masked'}} + return [replacements.get(id(message), message) for message in full] + + scanned = copy.deepcopy(subset) if copy_subset else subset + assert scanned == subset + result = call(redact, messages, scanned) + if result is not None and result is not messages: + state['messages'] = result + + assert messages[0] is subset[0] is first + assert state['messages'][1] is second + if copy_subset: + assert scanned[0] is not first + assert result is None + assert state['messages'] is messages + assert first['content'] == 'private' + elif redact_in_place: + assert result is state['messages'] is messages + assert first['content'] == 'masked' + else: + assert state['messages'] is result + assert result is not messages + assert result[0] is not first + assert result[0] == {'content': 'masked'} + assert first['content'] == 'private' +", + ) +} + +#[rstest] +#[case::adopt_replacement(false)] +#[case::merge_into_live_request(true)] +#[serial(python_interpreter)] +fn dispatcher_applies_returned_state_only_at_its_replacement_or_merge_boundary( + initialized_python: &InitializedPython, + #[case] merge: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + initialized_python.attach(|py| { + let globals = pattern_scope(py, backend, mode)?; + py.run( + c" +original = {'messages': [{'content': 'original'}], 'request_id': 'retained'} +replacement = {'messages': [{'content': 'redacted'}], 'verdict': 'allow'} +seen = [] + +def rewrite(data): + data['checkpoint'] = True + return replacement + +def observe(data): + seen.append(data) +", + Some(&globals), + None, + )?; + let original = item(&globals, "original"); + let old_messages = original.get_item("messages")?; + let replacement = item(&globals, "replacement"); + let call = item(&globals, "call"); + let returned = call.call1((item(&globals, "rewrite"), &original))?; + assert!(returned.is(&replacement)); + let current = if merge { + original.call_method1("update", (&returned,))?; + original.clone() + } else { + returned + }; + call.call1((item(&globals, "observe"), ¤t))?; + + assert!(item(&globals, "seen").get_item(0)?.is(¤t)); + assert!( + current + .get_item("messages")? + .is(replacement.get_item("messages")?) + ); + assert_eq!( + old_messages + .get_item(0)? + .get_item("content")? + .extract::()?, + "original" + ); + assert!(original.get_item("checkpoint")?.extract::()?); + assert_eq!(current.is(&original), merge); + if merge { + assert_eq!( + current.get_item("request_id")?.extract::()?, + "retained" + ); + assert!(current.get_item("checkpoint")?.extract::()?); + } else { + assert!(original.get_item("messages")?.is(&old_messages)); + assert!(!current.contains("request_id")?); + assert!(!current.contains("checkpoint")?); + } + Ok(()) + }) +} + +#[rstest] +#[case::failure_hook_only(false)] +#[case::success_hook_checks_block_flag(true)] +#[serial(python_interpreter)] +fn block_exception_preserves_stashed_context_for_later_hooks( + initialized_python: &InitializedPython, + #[case] dispatch_success: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| globals.set_item("dispatch_success", dispatch_success), + c" +def test_pattern(): + logger = Value() + logger.details = {} + request, failures, successes, success_calls = {}, [], [], [] + block = RuntimeError('blocked') + + def pre_call(data): + logger.details['blocked'] = True + data['logging_object'] = logger + raise block + + def success(details): + success_calls.append(details) + if not details.get('blocked'): + successes.append(details) + + def failure(data, error): + failures.append((data.pop('logging_object'), error)) + + with TestCase().assertRaises(RuntimeError) as caught: + call(pre_call, request) + assert caught.exception is block + assert request['logging_object'] is logger + assert logger.details == {'blocked': True} + + if dispatch_success: + call(success, logger.details) + assert successes == [] + assert success_calls[0] is logger.details + unblocked_details = {} + call(success, unblocked_details) + assert successes == [unblocked_details] + assert successes[0] is unblocked_details + assert len(success_calls) == (2 if dispatch_success else 0) + call(failure, request, caught.exception) + + assert len(failures) == 1 + assert failures[0][0] is logger + assert failures[0][1] is block + assert request == {} + assert logger.details == {'blocked': True} +", + ) +} + +#[rstest] +#[case::shared_inputs_lose_an_increment(false, false)] +#[case::reverse_completion_changes_writer_order(false, true)] +#[case::independent_snapshots_discard_edits(true, false)] +#[serial(python_interpreter)] +fn parallel_callbacks_share_state_without_isolation_or_automatic_return_merging( + initialized_python: &InitializedPython, + #[case] snapshots: bool, + #[case] reverse: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| { + globals.set_item("snapshots", snapshots)?; + globals.set_item("reverse", reverse) + }, + c" +async def test_pattern(): + state = {'count': 0, 'metadata': {'writers': []}} + entered = [asyncio.Event(), asyncio.Event()] + release = [asyncio.Event(), asyncio.Event()] + completed = [asyncio.Event(), asyncio.Event()] + + async def increment(data, index): + previous = data['count'] + entered[index].set() + await release[index].wait() + data['count'] = previous + 1 + data['metadata']['writers'].append(index) + completed[index].set() + return {'ignored_replacement': True} + + inputs = [copy.deepcopy(state) for _ in range(2)] if snapshots else [state, state] + tasks = [asyncio.create_task(invoke(increment, data, index)) for index, data in enumerate(inputs)] + await asyncio.gather(*(event.wait() for event in entered)) + assert state == {'count': 0, 'metadata': {'writers': []}} + assert all(not task.done() for task in tasks) + + order = [1, 0] if reverse else [0, 1] + for position, index in enumerate(order): + release[index].set() + await completed[index].wait() + assert inputs[index]['count'] == 1 + if position == 0: + assert not tasks[1 - index].done() + assert inputs[index]['metadata']['writers'] == [index] + + assert await asyncio.gather(*tasks) == [{'ignored_replacement': True}] * 2 + if snapshots: + assert state == {'count': 0, 'metadata': {'writers': []}} + assert inputs[0] is not inputs[1] + for index, data in enumerate(inputs): + assert data is not state + assert data == {'count': 1, 'metadata': {'writers': [index]}} + else: + assert inputs[0] is inputs[1] is state + assert state == {'count': 1, 'metadata': {'writers': order}} +", + ) +} + +#[rstest] +#[case::after_return(false)] +#[case::after_block_exception(true)] +#[serial(python_interpreter)] +fn background_task_keeps_arguments_and_updates_live_state_after_callback_finishes( + initialized_python: &InitializedPython, + #[case] raises: bool, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| globals.set_item("raises", raises), + c" +async def test_pattern(): + entered, release = asyncio.Event(), asyncio.Event() + tasks, observed_loops = [], [] + state = Value() + state.metadata = {'audit': 'pending'} + reference = weakref.ref(state) + event_loop = asyncio.get_running_loop() + + async def audit(data): + observed_loops.append(asyncio.get_running_loop()) + entered.set() + await release.wait() + data.metadata['audit'] = 'complete' + + def schedule(data): + observed_loops.append(asyncio.get_running_loop()) + tasks.append(asyncio.create_task(audit(data))) + if raises: + raise RuntimeError('blocked after scheduling') + + if raises: + with TestCase().assertRaisesRegex(RuntimeError, 'blocked after scheduling'): + await async_call(schedule, state) + else: + assert await async_call(schedule, state) is None + + metadata = state.metadata + snapshot = copy.deepcopy(metadata) + del state + await entered.wait() + gc.collect() + assert reference() is not None + assert metadata == {'audit': 'pending'} + assert len(tasks) == 1 + assert not tasks[0].done(), 'callback completion must not wait for background work' + assert tasks[0].get_loop() is event_loop + assert len(observed_loops) == 2 + assert all(loop is event_loop for loop in observed_loops) + + release.set() + assert await tasks[0] is None + assert metadata == {'audit': 'complete'} + assert snapshot == {'audit': 'pending'} + tasks.clear() + gc.collect() + assert reference() is None +", + ) +} + +#[rstest] +#[case::no_redaction_returns_original("passthrough")] +#[case::shallow_redaction_copies_only_response("shallow_redaction")] +#[case::per_key_copy_failure_shares_only_uncopyable_value("per_key_fallback")] +#[case::whole_copy_success_isolates_nested_values("deep_copy")] +#[case::whole_copy_failure_returns_original("whole_object_fallback")] +#[serial(python_interpreter)] +fn explicit_copy_boundaries_preserve_copy_depth_and_failure_fallback( + initialized_python: &InitializedPython, + #[case] copy_policy: &str, + #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, + #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode, +) -> PyResult<()> { + run_pattern( + initialized_python, + backend, + mode, + |globals| globals.set_item("copy_policy", copy_policy), + c" +def test_pattern(): + lock = threading.Lock() + state = {'messages': [{'content': 'private'}], 'response': {'content': 'private'}} + if copy_policy in ('per_key_fallback', 'whole_object_fallback'): + state['opaque'] = {'lock': lock, 'edits': []} + + def passthrough(data): + return data + + def shallow_redaction(data): + result = {**data, 'response': copy.deepcopy(data['response'])} + result['response']['content'] = 'masked' + return result + + def per_key_fallback(data): + result = {} + for key, value in data.items(): + try: + result[key] = copy.deepcopy(value) + except TypeError: + result[key] = value + return result + + def whole_object_fallback(data): + try: + return copy.deepcopy(data) + except TypeError: + return data + + callbacks = { + 'passthrough': passthrough, + 'shallow_redaction': shallow_redaction, + 'per_key_fallback': per_key_fallback, + 'deep_copy': whole_object_fallback, + 'whole_object_fallback': whole_object_fallback, + } + result = call(callbacks[copy_policy], state) + returns_original = copy_policy in ('passthrough', 'whole_object_fallback') + shares_messages = returns_original or copy_policy == 'shallow_redaction' + + assert (result is state) == returns_original + assert (result['response'] is state['response']) == returns_original + assert (result['messages'] is state['messages']) == shares_messages + result['messages'][0]['content'] = 'later edit' + assert state['messages'][0]['content'] == ('later edit' if shares_messages else 'private') + assert state['response']['content'] == 'private' + assert result['response']['content'] == ('masked' if copy_policy == 'shallow_redaction' else 'private') + + if 'opaque' in state: + assert result['opaque'] is state['opaque'] + assert result['opaque']['lock'] is lock + result['opaque']['edits'].append('shared despite copy') + assert state['opaque']['edits'] == ['shared despite copy'] +", + ) +} diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py index 6c123b64ca5..c841f444a13 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py @@ -1,5 +1,4 @@ import asyncio -import copy import gc import inspect import weakref @@ -31,46 +30,10 @@ class LiveCallFactory(CallFactory, Protocol): def live(self) -> int: ... -Arguments = tuple[tuple[object, ...], dict[str, object] | None] -ArgumentTransform = Callable[[tuple[object, ...], dict[str, object] | None], Arguments] - - -def reconstruct_envelope(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments: - return tuple(value for value in positional), None if keywords is None else dict(keywords) - - -def shallow_selected_payload(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments: - return (copy.copy(positional[0]), *positional[1:]), keywords - - -def deepcopy_graph(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments: - return copy.deepcopy((positional, keywords)) - - -def deepcopy_separate(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments: - return copy.deepcopy(positional), copy.deepcopy(keywords) - - def unchanged_result(value: object) -> object: return value -@dataclass(frozen=True, slots=True) -class ArgumentTransformFactory: - inner: CallFactory - transform: ArgumentTransform - - def prepare( - self, - callable: Callable[..., object], - positional: tuple[object, ...], - keywords: dict[str, object] | None = None, - awaited: bool = False, - ) -> PreparedInvocation: - args, kwargs = self.transform(positional, keywords) - return self.inner.prepare(callable, args, kwargs, awaited=awaited) - - @dataclass(frozen=True, slots=True) class ResultTransformInvocation: inner: PreparedInvocation @@ -209,22 +172,7 @@ def control_factory(control: str, inner: CallFactory) -> CallFactory: return CheckedWeakFactory() if control == "missing_handoff": return MissingHandoffFactory(inner) - if control in ("result_passthrough", "result_shallow", "result_deep"): - return ResultTransformFactory( - inner, - {"result_passthrough": unchanged_result, "result_shallow": copy.copy, "result_deep": copy.deepcopy}[ - control - ], - ) - return ArgumentTransformFactory( - inner, - { - "envelope": reconstruct_envelope, - "shallow_payload": shallow_selected_payload, - "deep_graph": deepcopy_graph, - "deep_separate": deepcopy_separate, - }[control], - ) + return {"result_passthrough": ResultTransformFactory(inner, unchanged_result)}[control] @dataclass @@ -232,82 +180,6 @@ class ControlNode: stage: int = 0 -@dataclass -class ControlPayload: - nested: ControlNode - stage: int = 0 - - -@dataclass(frozen=True, slots=True) -class IdentityObservation: - root: bool - nested: bool - cross_argument: bool - - -async def argument_identity(owners: CallFactory, awaited: bool) -> IdentityObservation: - nested = ControlNode() - original = ControlPayload(nested) - - def observe(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation: - return IdentityObservation(value is original, value.nested is nested, value.nested is alias) - - async def observe_async(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation: - return observe(value, alias=alias) - - owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited) - try: - pending = owner.invoke() - return await settle(pending, awaited) - finally: - owner.close() - - -@dataclass(frozen=True, slots=True) -class TimingObservation: - root: int - nested: int - alias: int - - -async def mutation_timing(owners: CallFactory, awaited: bool) -> TimingObservation: - nested = ControlNode() - original = ControlPayload(nested) - - def observe(value: ControlPayload, *, alias: ControlNode) -> TimingObservation: - return TimingObservation(value.stage, value.nested.stage, alias.stage) - - async def observe_async(value: ControlPayload, *, alias: ControlNode) -> TimingObservation: - return observe(value, alias=alias) - - owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited) - try: - original.stage = nested.stage = 1 - pending = owner.invoke() - original.stage = nested.stage = 2 - return await settle(pending, awaited) - finally: - owner.close() - - -async def result_identity(owners: CallFactory, awaited: bool) -> IdentityObservation: - original = ControlPayload(ControlNode()) - - def callback() -> ControlPayload: - return original - - async def callback_async() -> ControlPayload: - return original - - owner = owners.prepare(callback_async if awaited else callback, (), awaited=awaited) - try: - pending = owner.invoke() - result = await settle(pending, awaited) - return IdentityObservation(result is original, result.nested is original.nested, True) - finally: - owner.close() - - @dataclass class LifetimeCallback: awaited: bool @@ -444,19 +316,6 @@ async def direct_coroutine(owners: CallFactory, awaited: bool) -> bool: def expected_control(witness: str, control: str, awaited: bool) -> object: - if witness == "argument_identity": - return IdentityObservation( - control in ("identity", "envelope"), - control in ("identity", "envelope", "shallow_payload"), - control != "deep_separate", - ) - if witness == "mutation_timing": - stage = 2 if awaited else 1 - if control in ("deep_graph", "deep_separate"): - return TimingObservation(0, 0, 0) - return TimingObservation(0 if control == "shallow_payload" else stage, stage, stage) - if witness == "result_identity": - return IdentityObservation(control in ("identity", "result_passthrough"), control != "result_deep", True) if witness in ("deferred_lifetime", "pending_handoff"): if control == "weak" or (witness == "pending_handoff" and control == "missing_handoff"): return LifetimeObservation( @@ -480,9 +339,6 @@ def run_control(witness: str, control: str, retained: bool, awaited: bool, facto WITNESSES: dict[str, Callable[[CallFactory, bool], object]] = { - "argument_identity": argument_identity, - "mutation_timing": mutation_timing, - "result_identity": result_identity, "deferred_lifetime": deferred_lifetime, "borrowed_lifetime": borrowed_lifetime, "pending_handoff": pending_handoff, diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py index 430d7a656ea..a085155b572 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py @@ -10,6 +10,7 @@ from typing import Literal from unittest import TestCase import httpx +from fastapi import HTTPException import litellm from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -35,6 +36,12 @@ async def integration_invoke(owners, callback, *args, **kwargs): return await pending +async def integration_checkpoint(): + ready = asyncio.Event() + asyncio.get_running_loop().call_soon(ready.set) + await ready.wait() + + def integration_response(url, body, status=200, headers=None): return httpx.Response(status, json=body, headers=headers, request=httpx.Request("POST", url)) @@ -404,11 +411,18 @@ async def real_parallel_guardrail_snapshots(owners): for safe_memory_mode in (False, True): litellm.safe_memory_mode = safe_memory_mode await integration_parallel_snapshot_case(owners) + for ordinary_first in (False, True): + for reverse_completion in (False, True): + await integration_parallel_snapshot_case( + owners, ordinary_first=ordinary_first, reverse_completion=reverse_completion + ) + with TestCase().assertRaisesRegex(AssertionError, "copied live graph lost caller-visible mutation"): + await integration_parallel_snapshot_case(owners, copy_live=True) finally: litellm.safe_memory_mode = original_mode -async def integration_parallel_snapshot_case(owners): +async def integration_parallel_snapshot_case(owners, *, ordinary_first=None, reverse_completion=False, copy_live=False): from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.core_helpers import independent_snapshot from litellm.proxy.utils import ProxyLogging @@ -430,6 +444,21 @@ async def integration_parallel_snapshot_case(owners): assert raw is not live and raw["messages"][0] is not live["messages"][0] assert raw["uncopyable"] is sentinel and sentinel.attempts == 1 live["messages"][0]["content"] = "masked" + dispatched = independent_snapshot(live) if copy_live else live + ordinary = RuntimeError("ordinary guardrail failure") + blocking = HTTPException(status_code=400, detail={"error": "blocked by policy"}) + passthrough = ModifyResponseException("synthetic response", "fixture-model", dispatched) + errors = ( + {} + if ordinary_first is None + else { + "fixture-live-writer": passthrough, + "fixture-raw-a": ordinary if ordinary_first else blocking, + "fixture-raw-b": blocking if ordinary_first else ordinary, + } + ) + finished = [] + permits, completed = {}, {} class Inspect(CustomGuardrail): async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): @@ -447,7 +476,12 @@ async def integration_parallel_snapshot_case(owners): assert data["uncopyable"] is sentinel data["uncopyable"].observed.append(self.guardrail_name) else: - assert data is live and data["messages"][0]["content"] == "shared mutation" + assert data is dispatched and data["messages"][0]["content"] == "shared mutation" + await permits[self.guardrail_name].wait() + finished.append(self.guardrail_name) + completed[self.guardrail_name].set() + if self.guardrail_name in errors: + raise errors[self.guardrail_name] return {"discarded": self.guardrail_name} guardrails = tuple( @@ -462,30 +496,63 @@ async def integration_parallel_snapshot_case(owners): ) proxy = ProxyLogging.__new__(ProxyLogging) proxy.call_details = {"user_api_key_cache": DualCache()} + permits.update((guardrail.guardrail_name, asyncio.Event()) for guardrail in guardrails) + completed.update((guardrail.guardrail_name, asyncio.Event()) for guardrail in guardrails) task = asyncio.create_task( integration_invoke( - owners, proxy._run_parallel_pre_call_guardrails, guardrails, live, raw, UserAPIKeyAuth(), "acompletion" + owners, + proxy._run_parallel_pre_call_guardrails, + guardrails, + dispatched, + raw, + UserAPIKeyAuth(), + "acompletion", ) ) try: await arrived.wait() assert not task.done() - assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is live + assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is dispatched first, second = observations["fixture-raw-a"], observations["fixture-raw-b"] assert first is not second and first is not raw and second is not raw assert first["messages"][0] is not second["messages"][0] assert first["messages"][0] is not raw["messages"][0] assert first["uncopyable"] is second["uncopyable"] is raw["uncopyable"] is sentinel - assert sentinel.attempts == 3 and not sentinel.observed + assert sentinel.attempts == 3 + copy_live and not sentinel.observed release.set() - assert await task is None + order = tuple(completed)[:: -1 if reverse_completion else 1] + for index, name in enumerate(order): + assert not task.done() + permits[name].set() + await completed[name].wait() + await integration_checkpoint() + assert finished == list(order[: index + 1]) + if errors: + expected = ordinary if ordinary_first else blocking + with TestCase().assertRaises(type(expected)) as caught: + await task + assert caught.exception is expected + assert passthrough.request_data is dispatched + assert blocking.status_code == 400 + assert blocking.detail == { + "error": "blocked by policy", + "guardrail_name": "fixture-raw-b" if ordinary_first else "fixture-raw-a", + "guardrail_mode": GuardrailEventHooks.pre_call, + } + else: + assert await task is None assert raw["messages"] == [{"role": "user", "content": "original"}] - assert live["messages"][0]["content"] == "shared mutation" + assert dispatched["messages"][0]["content"] == "shared mutation" assert set(sentinel.observed) == {"fixture-raw-a", "fixture-raw-b"} and len(sentinel.observed) == 2 - assert "discarded" not in live + assert "discarded" not in dispatched assert first["messages"][0]["content"] == "fixture-raw-a" assert second["messages"][0]["content"] == "fixture-raw-b" - assert all(guardrail._pre_call_hook_already_ran(live) for guardrail in guardrails if guardrail.scan_raw_request) + assert all( + guardrail._pre_call_hook_already_ran(dispatched) is (guardrail.guardrail_name not in errors) + for guardrail in guardrails + if guardrail.scan_raw_request + ) + assert live["messages"][0]["content"] == "shared mutation", "copied live graph lost caller-visible mutation" finally: release.set() mutated.set() @@ -495,7 +562,13 @@ async def integration_parallel_snapshot_case(owners): async def real_purview_sync_background(owners): + for active_loop in (False, True): + await integration_purview_logging_case(owners, active_loop=active_loop) + + +async def integration_purview_logging_case(owners, *, active_loop): entered, release = threading.Event(), threading.Event() + async_entered, async_release = asyncio.Event(), asyncio.Event() calls, workers = [], [] main_thread = threading.get_ident() @@ -505,7 +578,11 @@ async def real_purview_sync_background(owners): calls.append((url, kwargs)) if url.endswith("/token"): entered.set() - assert release.wait(5), "background audit was not released" + if active_loop: + async_entered.set() + await async_release.wait() + else: + assert release.wait(5), "background audit was not released" return integration_response(url, {"access_token": "fixture-token", "expires_in": 3600}) assert kwargs["headers"]["Authorization"] == "Bearer fixture-token" if url.endswith("/compute"): @@ -535,19 +612,39 @@ async def real_purview_sync_background(owners): } result = ModelResponse(model="fixture-model", choices=[{"message": {"role": "assistant", "content": "before"}}]) owner = owners.prepare(purview.logging_hook, (kwargs, result, "completion"), awaited=False) + task = None try: - returned = await asyncio.to_thread(owner.invoke) + tasks_before = asyncio.all_tasks() + threads_before = set(threading.enumerate()) + returned = owner.invoke() if active_loop else await asyncio.to_thread(owner.invoke) owner.close() assert returned[0] is kwargs and returned[1] is result - assert await asyncio.to_thread(entered.wait, 5) - assert len(calls) == 1 and workers[0].ident != main_thread and workers[0].daemon + if active_loop: + await integration_checkpoint() + assert not calls and not workers and not entered.is_set() + assert asyncio.all_tasks() == tasks_before and set(threading.enumerate()) == threads_before + task = asyncio.create_task( + integration_invoke(owners, purview.async_logging_hook, kwargs, result, "completion") + ) + await async_entered.wait() + assert not task.done() and workers[0].ident == main_thread + else: + assert await asyncio.to_thread(entered.wait, 5) + assert workers[0].ident != main_thread and workers[0].daemon + assert len(calls) == 1 assert workers[0].is_alive() kwargs["messages"][0]["content"] = "too late for prompt extraction" kwargs["litellm_call_id"] = "after-http" result.choices[0].message.content = "response mutated while audit waits" release.set() - await asyncio.to_thread(workers[0].join, 5) - assert not workers[0].is_alive() and all(worker is workers[0] for worker in workers) + async_release.set() + if active_loop: + audited = await task + assert audited[0] is kwargs and audited[1] is result + else: + await asyncio.to_thread(workers[0].join, 5) + assert not workers[0].is_alive() + assert all(worker is workers[0] for worker in workers) assert len(calls) == 4 entries = [call[1]["json"]["contentToProcess"] for call in calls[2:]] assert [entry["activityMetadata"]["activity"] for entry in entries] == ["uploadText", "downloadText"] @@ -561,5 +658,10 @@ async def real_purview_sync_background(owners): finally: owner.close() release.set() - if workers: + async_release.set() + if task is not None: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + if workers and not active_loop: await asyncio.to_thread(workers[0].join, 5) diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py index f046595158c..118ca5f3013 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py @@ -563,58 +563,6 @@ async def repeated_ownership(owners): assert all(ref() is None for ref in refs) -async def retained_field_replacement(owners): - original = {"messages": [{"content": "original"}]} - replacement = {"messages": [{"content": "replacement"}]} - event = {"payload": original, "alias": original} - saved = [] - - def retain(value): - saved.append(value["payload"]) - - def replace(value): - value["payload"] = replacement - value["alias"]["messages"][0]["content"] = "mutated original" - - for callback in (retain, replace): - owner = owners.prepare(callback, (event,)) - try: - assert owner.invoke() is None - finally: - owner.close() - assert saved[0] is original is event["alias"] - assert event["payload"] is replacement - assert saved[0]["messages"][0]["content"] == "mutated original" - replacement["messages"][0]["content"] = "mutated replacement" - assert event["payload"]["messages"][0]["content"] == "mutated replacement" - assert original["messages"][0]["content"] == "mutated original" - - -async def queued_graph_ownership(owners): - queue = asyncio.Queue() - sentinel = Value() - reference = weakref.ref(sentinel) - payload = {"sentinel": sentinel, "nested": {"status": "queued"}} - snapshot = json.dumps(payload["nested"]) - enqueue = owners.prepare(queue.put_nowait, (payload,)) - try: - enqueue.invoke() - finally: - enqueue.close() - del sentinel, payload - gc.collect() - assert owners.live == 0 and reference() is not None - queued = queue.get_nowait() - queued["nested"]["status"] = "changed before flush" - assert json.loads(json.dumps(queued["nested"])) == {"status": "changed before flush"} - assert json.loads(snapshot) == {"status": "queued"} - assert queued["sentinel"] is reference() - queue.task_done() - del queued - gc.collect() - assert reference() is None - - async def detached_work_after_error(owners): for raises in (False, True): entered, release = asyncio.Event(), asyncio.Event() @@ -665,16 +613,18 @@ def run_checked(owners, scenario): async def run(): asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context)) - await asyncio.wait_for(scenario, timeout=15) + result = await asyncio.wait_for(scenario, timeout=15) gc.collect() assert owners.live == baseline pending = asyncio.all_tasks() - {asyncio.current_task()} assert not pending, f"undrained tasks: {pending}" + return result - asyncio.run(run()) + result = asyncio.run(run()) gc.collect() assert not background_failures, f"unhandled background failures: {background_failures}" assert owners.live == baseline + return result def run_scenario(name, retained, factory): diff --git a/litellm-rust/crates/python-interop/tests/prepared_call.rs b/litellm-rust/crates/python-interop/tests/prepared_call.rs index d5d444d31b6..6ddad3f2510 100644 --- a/litellm-rust/crates/python-interop/tests/prepared_call.rs +++ b/litellm-rust/crates/python-interop/tests/prepared_call.rs @@ -1,6 +1,7 @@ use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; +use pyo3::exceptions::{PyAssertionError, PyKeyboardInterrupt, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; +use pyo3::types::{PyDict, PyList, PyTuple}; use rstest::rstest; #[path = "support/mod.rs"] @@ -17,19 +18,23 @@ fn retains_aliases_mutations_and_original_result( let globals = scope( py, c" -shared = {'value': 'before'} -payload = {'nested': shared} -saved = [] def callback(data, *, alias): - assert data['nested'] is alias + observed.append(data['nested'] is alias) saved.append(data) alias['value'] = 'during' return data ", )?; - let payload = item(&globals, "payload"); + let shared = PyDict::new(py); + shared.set_item("value", "before")?; + let payload = PyDict::new(py); + payload.set_item("nested", &shared)?; + let saved = PyList::empty(py); + let observed = PyList::empty(py); + globals.set_item("saved", &saved)?; + globals.set_item("observed", &observed)?; let keywords = PyDict::new(py); - keywords.set_item("alias", item(&globals, "shared"))?; + keywords.set_item("alias", &shared)?; let invocation = PreparedCall::new( InvocationMode::Direct, item(&globals, "callback").unbind(), @@ -39,16 +44,19 @@ def callback(data, *, alias): let result = invoke_direct(&invocation, py)?; assert!(result.bind(py).is(&payload)); drop(invocation); - py.run( - c" -assert saved[0] is payload -assert shared['value'] == 'during' -shared['value'] = 'after' -assert saved[0]['nested']['value'] == 'after' -", - Some(&globals), - None, - ) + assert_eq!(observed.extract::>()?, [true]); + assert!(saved.get_item(0)?.is(&payload)); + assert_eq!(item(&shared, "value").extract::()?, "during"); + shared.set_item("value", "after")?; + assert_eq!( + saved + .get_item(0)? + .get_item("nested")? + .get_item("value")? + .extract::()?, + "after" + ); + Ok(()) }) } @@ -61,33 +69,38 @@ fn preserves_exception_identity_cause_traceback_and_prior_mutation( let globals = scope( py, c" -payload = {} -error = KeyboardInterrupt('original') -cause = ValueError('cause') def callback(data): data['changed'] = True raise error from cause ", )?; + let payload = PyDict::new(py); + let original = PyKeyboardInterrupt::new_err("original"); + let cause = PyValueError::new_err("cause"); + globals.set_item("error", original.value(py))?; + globals.set_item("cause", cause.value(py))?; let invocation = PreparedCall::new( InvocationMode::Direct, item(&globals, "callback").unbind(), - PyTuple::new(py, [item(&globals, "payload")])?.unbind(), + PyTuple::new(py, [&payload])?.unbind(), None, ); let error = invoke_direct(&invocation, py).unwrap_err(); - assert!(error.value(py).is(item(&globals, "error"))); + assert!(error.value(py).is(original.value(py))); drop(invocation); - py.run( - c" -import traceback -assert payload['changed'] is True -assert error.__cause__ is cause -assert traceback.extract_tb(error.__traceback__)[-1].name == 'callback' -", - Some(&globals), - None, - ) + assert!(item(&payload, "changed").extract::()?); + assert!(error.value(py).getattr("__cause__")?.is(cause.value(py))); + let frames = py + .import("traceback")? + .call_method1("extract_tb", (error.traceback(py),))?; + assert_eq!( + frames + .get_item(frames.len()? - 1)? + .getattr("name")? + .extract::()?, + "callback" + ); + Ok(()) }) } @@ -98,15 +111,16 @@ fn returns_coroutine_without_executing_it(initialized_python: &InitializedPython let globals = scope( py, c" -import inspect -started = [] async def work(): started.append(True) -coroutine = work() def callback(): return coroutine ", )?; + let started = PyList::empty(py); + globals.set_item("started", &started)?; + let coroutine = item(&globals, "work").call0()?; + globals.set_item("coroutine", &coroutine)?; let invocation = PreparedCall::new( InvocationMode::Direct, item(&globals, "callback").unbind(), @@ -114,16 +128,13 @@ def callback(): None, ); let result = invoke_direct(&invocation, py)?; - assert!(result.bind(py).is(item(&globals, "coroutine"))); - py.run( - c" -assert started == [] -assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CREATED -coroutine.close() -", - Some(&globals), - None, - ) + let inspect = py.import("inspect")?; + let state = inspect.call_method1("getcoroutinestate", (&coroutine,)); + coroutine.call_method0("close")?; + assert!(result.bind(py).is(&coroutine)); + assert!(started.is_empty()); + assert!(state?.eq(inspect.getattr("CORO_CREATED")?)?); + Ok(()) }) } @@ -149,25 +160,31 @@ fn preserves_current_context_thread_and_reentry( let globals = scope( py, c" -import contextvars import threading -context = contextvars.ContextVar('prepared_call_context') -token = context.set('caller') -thread = threading.get_ident() -payload = {} def inner(data): - assert context.get() == 'outer' - assert threading.get_ident() == thread + observed.append((context.get(), threading.get_ident())) data['inner'] = True context.set('inner') return data def outer(): - assert context.get() == 'caller' - assert threading.get_ident() == thread + observed.append((context.get(), threading.get_ident())) context.set('outer') return reenter(inner, payload) ", )?; + let context = py + .import("contextvars")? + .getattr("ContextVar")? + .call1(("prepared_call_context",))?; + let thread = py + .import("threading")? + .call_method0("get_ident")? + .extract::()?; + let payload = PyDict::new(py); + let observed = PyList::empty(py); + globals.set_item("context", &context)?; + globals.set_item("payload", &payload)?; + globals.set_item("observed", &observed)?; globals.set_item("reenter", wrap_pyfunction!(reenter, py)?)?; let invocation = PreparedCall::new( InvocationMode::Direct, @@ -175,19 +192,18 @@ def outer(): PyTuple::empty(py).unbind(), None, ); - let result = invoke_direct(&invocation, py)?; - assert!(result.bind(py).is(item(&globals, "payload"))); - py.run( - c" -try: - assert payload['inner'] is True - assert context.get() == 'inner' -finally: - context.reset(token) -", - Some(&globals), - None, - ) + let token = context.call_method1("set", ("caller",))?; + let result = invoke_direct(&invocation, py); + let final_context = context.call_method0("get"); + context.call_method1("reset", (token,))?; + assert!(result?.bind(py).is(&payload)); + assert!(item(&payload, "inner").extract::()?); + assert_eq!(final_context?.extract::()?, "inner"); + assert_eq!( + observed.extract::>()?, + [("caller".to_owned(), thread), ("outer".to_owned(), thread)] + ); + Ok(()) }) } @@ -200,57 +216,56 @@ fn owns_arguments_until_release_and_preserves_callback_retention( let globals = scope( py, c" -import gc -import weakref -saved = [] class Value: pass class Callback: def __call__(self, value, *, other): saved.append(value) - assert other is other_ref() -value = Value() -other = Value() -callback = Callback() -value_ref = weakref.ref(value) -other_ref = weakref.ref(other) -callback_ref = weakref.ref(callback) + observed.append(other is other_ref()) ", )?; + globals.set_item("saved", PyList::empty(py))?; + globals.set_item("observed", PyList::empty(py))?; + let value = item(&globals, "Value").call0()?; + let other = item(&globals, "Value").call0()?; + let callback = item(&globals, "Callback").call0()?; + let weakref = py.import("weakref")?; + for (name, object) in [ + ("value_ref", &value), + ("other_ref", &other), + ("callback_ref", &callback), + ] { + globals.set_item(name, weakref.call_method1("ref", (object,))?)?; + } let keywords = PyDict::new(py); - keywords.set_item("other", item(&globals, "other"))?; + keywords.set_item("other", other)?; let invocation = PreparedCall::new( InvocationMode::Direct, - item(&globals, "callback").unbind(), - PyTuple::new(py, [item(&globals, "value")])?.unbind(), + callback.unbind(), + PyTuple::new(py, [value])?.unbind(), Some(keywords.unbind()), ); - py.run(c"del value, other, callback", Some(&globals), None)?; Ok::<_, PyErr>((invocation, globals.unbind())) })?; Python::attach(|py| { let globals = globals.bind(py); - py.run( - c"assert all(ref() is not None for ref in (value_ref, other_ref, callback_ref))", - Some(globals), - None, - )?; + for name in ["value_ref", "other_ref", "callback_ref"] { + assert!(!item(globals, name).call0()?.is_none(), "{name}"); + } assert!(invoke_direct(&invocation, py)?.is_none(py)); drop(invocation); - py.run( - c" -gc.collect() -assert callback_ref() is None -assert other_ref() is None -assert value_ref() is saved[0] -saved[0].still_usable = True -saved.clear() -gc.collect() -assert value_ref() is None -", - Some(globals), - None, - ) + let gc = py.import("gc")?; + gc.call_method0("collect")?; + assert_eq!(item(globals, "observed").extract::>()?, [true]); + assert!(item(globals, "callback_ref").call0()?.is_none()); + assert!(item(globals, "other_ref").call0()?.is_none()); + let saved = item(globals, "saved"); + assert!(item(globals, "value_ref").call0()?.is(saved.get_item(0)?)); + saved.get_item(0)?.setattr("still_usable", true)?; + saved.call_method0("clear")?; + gc.call_method0("collect")?; + assert!(item(globals, "value_ref").call0()?.is_none()); + Ok(()) }) } @@ -292,31 +307,51 @@ fn checked_runner_rejects_unhandled_background_failures( async def fail(): raise RuntimeError('background task regression') -for cyclic in (False, True): - for handled in (False, True): - async def scenario(cyclic=cyclic, handled=handled): - task = asyncio.create_task(fail()) - if cyclic: - task.cycle = task - await checkpoint() - assert task.done() - if handled: - with TestCase().assertRaisesRegex(RuntimeError, 'background task regression'): - task.result() - del task - - owners = ReferenceFactory() - if handled: - run_checked(owners, scenario()) - else: - with TestCase().assertRaisesRegex( - AssertionError, r'unhandled background failures: .*background task regression' - ): - run_checked(owners, scenario()) +async def scenario(cyclic, handled, observed): + task = asyncio.create_task(fail()) + if cyclic: + task.cycle = task + await checkpoint() + observed['done'] = task.done() + if handled: + try: + task.result() + except RuntimeError as error: + observed['error'] = str(error) + del task ", Some(&globals), None, - ) + )?; + for cyclic in [false, true] { + for handled in [false, true] { + let observed = PyDict::new(py); + let owners = item(&globals, "ReferenceFactory").call0()?; + let scenario = item(&globals, "scenario").call1((cyclic, handled, &observed))?; + let result = item(&globals, "run_checked").call1((owners, scenario)); + assert!( + item(&observed, "done").extract::()?, + "cyclic={cyclic}, handled={handled}" + ); + if handled { + result?; + assert_eq!( + item(&observed, "error").extract::()?, + "background task regression" + ); + } else { + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py), "{error}"); + let message = error.value(py).str()?.to_str()?.to_owned(); + assert!( + message.starts_with("unhandled background failures: "), + "{message}" + ); + assert!(message.contains("background task regression"), "{message}"); + } + } + } + Ok(()) }) } @@ -330,11 +365,7 @@ fn real_ocr_logging_preserves_execution_roots_and_continues_after_error( let globals = scope( py, c" -from datetime import datetime from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.litellm_logging import Logging - -order = [] class Retain(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): @@ -365,60 +396,107 @@ class Observe(CustomLogger): 'document' in self.view['complete_input_dict'], ) -first = Retain() -last = Observe() -document = {'value': 'original'} -headers = {'X-Trace': 'original'} -body = {'document': document, 'alias': document} -view = {'headers': headers, 'complete_input_dict': body, 'api_base': 'https://example.invalid/ocr'} -logger = Logging( - model='test', messages=[], stream=False, call_type='ocr', - start_time=datetime.now(), litellm_call_id='retained-test', function_id='retained-test', - dynamic_input_callbacks=[first, MutateThenFail(), last], -) ", )?; - let headers = item(&globals, "headers").unbind(); - let body = item(&globals, "body").unbind(); - let view = item(&globals, "view").cast_into::()?.unbind(); - let invocation = prepare_pre_call(py, &item(&globals, "logger"), view.bind(py))?; + let order = PyList::empty(py); + globals.set_item("order", &order)?; + let first = item(&globals, "Retain").call0()?; + let last = item(&globals, "Observe").call0()?; + let document = PyDict::new(py); + document.set_item("value", "original")?; + let headers = PyDict::new(py); + headers.set_item("X-Trace", "original")?; + let body = PyDict::new(py); + body.set_item("document", &document)?; + body.set_item("alias", &document)?; + let view = PyDict::new(py); + view.set_item("headers", &headers)?; + view.set_item("complete_input_dict", &body)?; + view.set_item("api_base", "https://example.invalid/ocr")?; + let keywords = PyDict::new(py); + keywords.set_item("model", "test")?; + keywords.set_item("messages", PyList::empty(py))?; + keywords.set_item("stream", false)?; + keywords.set_item("call_type", "ocr")?; + keywords.set_item( + "start_time", + py.import("datetime")? + .getattr("datetime")? + .call_method0("now")?, + )?; + keywords.set_item("litellm_call_id", "retained-test")?; + keywords.set_item("function_id", "retained-test")?; + keywords.set_item( + "dynamic_input_callbacks", + PyList::new( + py, + [&first, &item(&globals, "MutateThenFail").call0()?, &last], + )?, + )?; + let logger = py + .import("litellm.litellm_core_utils.litellm_logging")? + .getattr("Logging")? + .call((), Some(&keywords))?; + let invocation = prepare_pre_call(py, &logger, &view)?; assert!(invoke_direct(&invocation, py)?.is_none(py)); drop(invocation); - py.run(c"del headers, body, view", Some(&globals), None)?; - assert!( - headers - .bind(py) - .is(item(&globals, "first").getattr("headers")?) - ); - assert!(body.bind(py).is(item(&globals, "first").getattr("body")?)); - assert!(view.bind(py).is(item(&globals, "last").getattr("view")?)); + assert!(headers.is(first.getattr("headers")?)); + assert!(body.is(first.getattr("body")?)); + assert!(view.is(last.getattr("view")?)); + assert_eq!(item(&headers, "X-Trace").extract::()?, "mutated"); assert_eq!( - headers.bind(py).get_item("X-Trace")?.extract::()?, + order.extract::>()?, + ["retain", "mutate_then_fail", "observe"] + ); + assert_eq!( + first.getattr("snapshot")?.extract::<(String, String)>()?, + ("original".to_owned(), "original".to_owned()) + ); + assert_eq!( + last.getattr("snapshot")? + .extract::<(Vec<(String, String)>, bool, bool)>()?, + ( + vec![("X-Trace".to_owned(), "replacement".to_owned())], + true, + false + ) + ); + assert!(first.getattr("view")?.is(last.getattr("view")?)); + assert!(first.getattr("body")?.get_item("document")?.is(&document)); + assert!(first.getattr("body")?.get_item("alias")?.is(&document)); + assert_eq!(item(&document, "value").extract::()?, "mutated"); + assert_eq!( + last.getattr("view")? + .get_item("headers")? + .get_item("X-Trace")? + .extract::()?, + "replacement" + ); + let replacement = PyDict::new(py); + replacement.set_item("replacement", true)?; + assert!( + last.getattr("view")? + .get_item("complete_input_dict")? + .eq(replacement)? + ); + document.set_item("value", "after invocation")?; + assert_eq!( + first + .getattr("body")? + .get_item("document")? + .get_item("value")? + .extract::()?, + "after invocation" + ); + drop((headers, body, view)); + assert_eq!( + first + .getattr("headers")? + .get_item("X-Trace")? + .extract::()?, "mutated" ); - py.run( - c" -assert order == ['retain', 'mutate_then_fail', 'observe'] -assert first.snapshot == ('original', 'original') -assert last.snapshot == ((('X-Trace', 'replacement'),), True, False) -assert first.view is last.view -assert first.body['document'] is document -assert first.body['alias'] is document -assert document['value'] == 'mutated' -assert last.view['headers']['X-Trace'] == 'replacement' -assert last.view['complete_input_dict'] == {'replacement': True} -document['value'] = 'after invocation' -assert first.body['document']['value'] == 'after invocation' -", - Some(&globals), - None, - )?; - drop((headers, body, view)); - py.run( - c"assert first.headers['X-Trace'] == 'mutated'", - Some(&globals), - None, - ) + Ok(()) }) } diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 8ecb2724acb..f80e6252b59 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -28,8 +28,6 @@ from litellm.llms.base_llm.ocr.transformation import ( parse_ocr_request_format, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -242,19 +240,6 @@ async def aocr( ) ``` """ - if rust_enabled(): - return await rust_ocr_bridge.aocr( - { - **kwargs, - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - } - ) return await _legacy_aocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) @@ -525,20 +510,6 @@ def ocr( print(f"Page {page.index}: {page.markdown}") ``` """ - if rust_enabled(): - arguments: Final[dict[str, object]] = { - **kwargs, - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - } - if kwargs.get("aocr") is True: - return rust_ocr_bridge.aocr(arguments) - return rust_ocr_bridge.ocr(arguments) return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 9ced40eab8d..908b079e53d 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,7 +2,6 @@ from __future__ import annotations -import inspect import traceback from collections.abc import Awaitable from contextvars import copy_context @@ -58,14 +57,67 @@ async def aocr(arguments: dict[str, object]) -> OCRResponse: def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> object: import litellm - from litellm.litellm_core_utils.litellm_logging import Logging + from litellm import utils + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils import litellm_logging + from litellm.litellm_core_utils.coroutine_checker import coroutine_checker + from litellm.litellm_core_utils.litellm_logging import Logging, set_callbacks supplied: Final = arguments.get("litellm_logging_obj") if supplied is not None: return supplied - callbacks: Final = tuple(dict.fromkeys((*litellm.callbacks, *cast(list, arguments.get("callbacks") or [])))) + callbacks: Final = tuple(dict.fromkeys(utils.get_dynamic_callbacks(cast(list, arguments.get("callbacks"))))) success: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("success_callback") or [])))) failure: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("failure_callback") or [])))) + configured: Final = tuple( + dict.fromkeys( + ( + *litellm.input_callback, + *litellm.success_callback, + *litellm.failure_callback, + *litellm._async_success_callback, + *litellm._async_failure_callback, + *success, + *failure, + ) + ) + ) + uninitialized: Final = [ + cb + for cb in configured + if isinstance(cb, str) + and ( + cb not in litellm._known_custom_logger_compatible_callbacks + or cb in litellm.input_callback + litellm.success_callback + litellm.failure_callback + ) + and cb not in (utils.callback_list or []) + ] + if uninitialized: + set_callbacks(uninitialized, function_id=arguments.get("id")) + utils.callback_list = list(dict.fromkeys((*(utils.callback_list or []), *uninitialized))) + if litellm_logging.customLogger is None: + set_callbacks([cb for cb in configured if callable(cb)], function_id=arguments.get("id")) + for event, registered, add_async in ( + ("input", litellm.input_callback, litellm.logging_callback_manager.add_litellm_input_callback), + ("success", litellm.success_callback, litellm.logging_callback_manager.add_litellm_async_success_callback), + ("failure", litellm.failure_callback, litellm.logging_callback_manager.add_litellm_async_failure_callback), + ): + for cb in tuple(registered): + if coroutine_checker.is_async_callable(cb) or (event == "success" and cb in ("dynamodb", "openmeter")): + if cb not in getattr(litellm, f"_async_{event}_callback"): + add_async(cb) + registered.remove(cb) + elif event != "input" and isinstance(cb, str) and cb in litellm._known_custom_logger_compatible_callbacks: + utils._add_custom_logger_callback_to_specific_event(cb, event) + for event, registered, add_sync in ( + ("success", litellm._async_success_callback, litellm.logging_callback_manager.add_litellm_success_callback), + ("failure", litellm._async_failure_callback, litellm.logging_callback_manager.add_litellm_failure_callback), + ): + for cb in tuple(registered): + if callable(cb) and not isinstance(cb, CustomLogger) and not coroutine_checker.is_async_callable(cb): + if cb not in getattr(litellm, f"{event}_callback"): + add_sync(cb) + registered.remove(cb) call_id: Final = str(arguments.get("litellm_call_id") or uuid4()) logger: Final = Logging( model=str(arguments["model"]), @@ -76,14 +128,27 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> obje litellm_call_id=call_id, function_id=str(arguments.get("id") or ""), litellm_trace_id=cast(str | None, arguments.get("litellm_trace_id")), - dynamic_input_callbacks=[cb for cb in callbacks if cb not in litellm.input_callback], - dynamic_success_callbacks=[cb for cb in success if not inspect.iscoroutinefunction(cb)], - dynamic_async_success_callbacks=list(success), - dynamic_failure_callbacks=[cb for cb in failure if not inspect.iscoroutinefunction(cb)], - dynamic_async_failure_callbacks=list(failure), + dynamic_input_callbacks=[ + cb for cb in callbacks if cb not in litellm.input_callback and not coroutine_checker.is_async_callable(cb) + ], + dynamic_success_callbacks=[ + cb for cb in success if not coroutine_checker.is_async_callable(cb) and cb not in ("dynamodb", "s3") + ], + dynamic_async_success_callbacks=[ + cb + for cb in success + if coroutine_checker.is_async_callable(cb) or isinstance(cb, CustomLogger) or cb in ("dynamodb", "s3") + ], + dynamic_failure_callbacks=[cb for cb in failure if not coroutine_checker.is_async_callable(cb)], + dynamic_async_failure_callbacks=[ + cb for cb in failure if coroutine_checker.is_async_callable(cb) or isinstance(cb, CustomLogger) + ], kwargs=arguments, supports_correlation_logging=asynchronous, ) + logger.dynamic_input_callbacks = [ + cb for cb in dict.fromkeys(logger.dynamic_input_callbacks or []) if cb not in litellm.input_callback + ] arguments["litellm_call_id"] = call_id arguments["litellm_logging_obj"] = logger return logger @@ -110,7 +175,13 @@ def invoke_terminal( _retained: Final = roots await logging.async_success_handler(value, start_time, end_time) - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=run_async()) + def enqueue() -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=run_async()) + + if getattr(logging, "_defer_async_logging", False) is True: + logging._enqueue_deferred_logging = enqueue + else: + enqueue() return None if action == "sync_success_if_needed": if logging._should_run_sync_callbacks_for_async_calls(): diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 806a1130c2d..8840ab9183c 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,4 +1,4 @@ -"""Strict whole-argument OCR dispatch and opt-in native transport contracts.""" +"""Public Python OCR routing and private native OCR proof contracts.""" import asyncio import atexit @@ -96,14 +96,7 @@ def no_python_ocr(monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("enable_with", ["global", "environment"]) -async def test_unwrapped_dispatch_preserves_every_argument( - monkeypatch, no_python_ocr, injected_native, response, asynchronous, enable_with -): - if enable_with == "global": - litellm.rust(True) - else: - monkeypatch.setenv("LITELLM_RUST", "1") +async def test_bridge_preserves_every_argument(no_python_ocr, injected_native, response, asynchronous): opaque = object() document = {"type": "file", "file": opaque, "mime_type": "application/pdf"} metadata = {"opaque": opaque, "nested": []} @@ -129,9 +122,7 @@ async def test_unwrapped_dispatch_preserves_every_argument( "kwargs": {"caller_owned": opaque}, "aocr": opaque, } - result = ( - await inspect.unwrap(ocr_main.aocr)(**arguments) if asynchronous else inspect.unwrap(ocr_main.ocr)(**arguments) - ) + result = await rust_bridge.aocr(arguments) if asynchronous else rust_bridge.ocr(arguments) sync, async_native = injected_native selected, unused = (async_native, sync) if asynchronous else (sync, async_native) @@ -152,34 +143,27 @@ async def test_unwrapped_dispatch_preserves_every_argument( @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_unwrapped_dispatch_keeps_unresolved_defaults( - no_python_ocr, injected_native, document, response, asynchronous -): +async def test_public_wrapper_passes_defaults_to_legacy(monkeypatch, injected_native, document, response, asynchronous): litellm.rust(True) - result = ( - await inspect.unwrap(ocr_main.aocr)(MODEL, document) - if asynchronous - else inspect.unwrap(ocr_main.ocr)(MODEL, document) - ) - selected = injected_native[int(asynchronous)] - selected.assert_called_once_with( - { - "model": MODEL, - "document": document, - "api_key": None, - "api_base": None, - "timeout": None, - "custom_llm_provider": None, - "extra_headers": None, - } - ) + legacy = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(ocr_main, "_legacy_aocr" if asynchronous else "_legacy_ocr", legacy) + result = await litellm.aocr(MODEL, document) if asynchronous else litellm.ocr(MODEL, document) + legacy.assert_called_once_with(MODEL, document, None, None, None, None, None) + if asynchronous: + legacy.assert_awaited_once() assert result is response + for native in injected_native: + native.assert_not_called() @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_public_decorator_routes_full_kwargs(no_python_ocr, injected_native, document, response, asynchronous): +async def test_public_wrapper_passes_full_kwargs_to_legacy( + monkeypatch, injected_native, document, response, asynchronous +): litellm.rust(True) + legacy = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(ocr_main, "_legacy_aocr" if asynchronous else "_legacy_ocr", legacy) metadata = {"test_tag": "whole-arguments"} pages = [0, 2] arguments = { @@ -197,21 +181,28 @@ async def test_public_decorator_routes_full_kwargs(no_python_ocr, injected_nativ } result = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - selected = injected_native[int(asynchronous)] - selected.assert_called_once() - injected_native[not asynchronous].assert_not_called() - (forwarded,) = selected.call_args.args + legacy.assert_called_once_with( + MODEL, + document, + "sk-test", + "https://example.invalid", + 12.5, + None, + arguments["extra_headers"], + pages=pages, + include_image_base64=True, + metadata=metadata, + arbitrary_option=arguments["arbitrary_option"], + num_retries=0, + ) + if asynchronous: + legacy.assert_awaited_once() assert result is response - for name in ("model", "api_key", "api_base", "timeout", "extra_headers", "arbitrary_option", "num_retries"): - assert forwarded[name] == arguments[name] - assert forwarded["document"] is document - assert forwarded["pages"] is pages - assert forwarded["include_image_base64"] is True - assert forwarded["metadata"]["test_tag"] == "whole-arguments" - assert forwarded["custom_llm_provider"] is None - assert "litellm_logging_obj" not in forwarded - assert "litellm_call_id" not in forwarded - assert "kwargs" not in forwarded + assert legacy.call_args.args[1] is document + assert legacy.call_args.kwargs["pages"] is pages + assert legacy.call_args.kwargs["metadata"] is metadata + for native in injected_native: + native.assert_not_called() @pytest.mark.asyncio @@ -227,10 +218,9 @@ async def test_bridge_passes_same_dictionary_and_response(injected_native, docum @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("public", [False, True], ids=["unwrapped", "decorated"]) @pytest.mark.parametrize("failure", ["missing", "unsupported", "runtime"]) -async def test_native_failures_propagate_without_fallback( - monkeypatch, no_python_ocr, injected_native, document, asynchronous, public, failure +async def test_bridge_failures_propagate_without_fallback( + monkeypatch, no_python_ocr, injected_native, document, asynchronous, failure ): litellm.rust(True) error = ( @@ -244,12 +234,9 @@ async def test_native_failures_propagate_without_fallback( monkeypatch.setattr(rust_bridge_bindings, "get_native_bridge", lambda: None) else: injected_native[int(asynchronous)].side_effect = error - function = litellm.aocr if asynchronous else litellm.ocr - route = function if public else inspect.unwrap(function) + arguments = dict(model=MODEL, document=document, api_key="sk-test", num_retries=0) with pytest.raises(RuntimeError if failure == "missing" else type(error)) as caught: - result = route(model=MODEL, document=document, api_key="sk-test", num_retries=0) - if asynchronous: - await result + await rust_bridge.aocr(arguments) if asynchronous else rust_bridge.ocr(arguments) if failure == "missing": assert "OCR" in str(caught.value).upper() else: @@ -263,26 +250,27 @@ async def test_native_failures_propagate_without_fallback( async def test_bridge_missing_binding_is_strict(document, asynchronous): rust_bridge._OCR.override(None) rust_bridge._AOCR.override(None) - with pytest.raises(RuntimeError, match="(?i)ocr"): - if asynchronous: - await rust_bridge.aocr({"model": MODEL, "document": document}) - else: - rust_bridge.ocr({"model": MODEL, "document": document}) + arguments = {"model": MODEL, "document": document} + with pytest.raises(RuntimeError, match=r"(?i)ocr"): + await rust_bridge.aocr(arguments) if asynchronous else rust_bridge.ocr(arguments) @pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("setting", ["default", "disabled", "overrides-environment"]) -async def test_rust_off_keeps_python_preparation_and_transport( - monkeypatch, injected_native, response, asynchronous, setting +@pytest.mark.parametrize("route", ["ocr", "aocr", "ocr-async"]) +@pytest.mark.parametrize("setting", ["default", "disabled", "overrides-environment", "global", "environment"]) +async def test_public_ocr_always_keeps_python_preparation_and_transport( + monkeypatch, injected_native, response, route, setting ): - if setting == "overrides-environment": + asynchronous = route != "ocr" + if setting in ("overrides-environment", "environment"): monkeypatch.setenv("LITELLM_RUST", "1") - if setting != "default": + if setting in ("disabled", "overrides-environment"): litellm.rust(False) + if setting == "global": + litellm.rust(True) prepare = Mock(wraps=ocr_main._prepare_ocr_request) handler = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) - lookup = Mock(side_effect=lambda: pytest.fail("disabled Rust binding was consulted")) + lookup = Mock(side_effect=lambda: pytest.fail("public OCR consulted a native binding")) monkeypatch.setattr(ocr_main, "_prepare_ocr_request", prepare) monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", handler) monkeypatch.setattr(rust_bridge, "load_rust_ocr", lookup) @@ -296,8 +284,10 @@ async def test_rust_off_keeps_python_preparation_and_transport( "pages": [0], "include_image_base64": True, "num_retries": 0, + **({"aocr": True} if route == "ocr-async" else {}), } - result = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + function = litellm.aocr if route == "aocr" else litellm.ocr + result = await function(**arguments) if asynchronous else function(**arguments) assert result is response prepare.assert_called_once() @@ -322,8 +312,9 @@ async def test_rust_off_keeps_python_preparation_and_transport( @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_rust_off_preserves_python_exception_mapping(monkeypatch, injected_native, document, asynchronous): - litellm.rust(False) +@pytest.mark.parametrize("enabled", [False, True]) +async def test_legacy_preserves_python_exception_mapping(monkeypatch, injected_native, document, asynchronous, enabled): + litellm.rust(enabled) original_error = ValueError("Python transport failed") mapped_error = RuntimeError("mapped Python error") mapping = Mock(return_value=mapped_error) @@ -332,10 +323,9 @@ async def test_rust_off_preserves_python_exception_mapping(monkeypatch, injected monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", handler) arguments = dict(model=MODEL, document=document, api_key="sk-test", litellm_logging_obj=Mock()) with pytest.raises(RuntimeError) as caught: - if asynchronous: - await inspect.unwrap(ocr_main._legacy_aocr)(**arguments) - else: - inspect.unwrap(ocr_main._legacy_ocr)(**arguments) + await inspect.unwrap(ocr_main._legacy_aocr)(**arguments) if asynchronous else inspect.unwrap( + ocr_main._legacy_ocr + )(**arguments) assert caught.value is mapped_error handler.assert_called_once() mapping.assert_called_once() @@ -393,7 +383,8 @@ def test_loader_caches_missing_extension_until_reset(monkeypatch): @pytest.fixture -def native_ocr(): +def native_ocr(monkeypatch, reset_rust_state): + """PRIVATE, test-only route selection; public OCR stays Python until full lifecycle parity.""" native = rust_bridge_loader.get_native_bridge() try: available = native is not None and all( @@ -406,9 +397,69 @@ def native_ocr(): if os.environ.get("LITELLM_REQUIRE_NATIVE_OCR") == "1": pytest.fail(message) pytest.skip(message) + python_ocr, python_aocr = litellm.ocr, litellm.aocr + signature = inspect.signature(python_ocr) + + def native_arguments(args, kwargs): + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + return {**bound.arguments.pop("kwargs"), **bound.arguments} + + def private_ocr(*args, **kwargs): + if not configuration.rust_enabled(): + return python_ocr(*args, **kwargs) + arguments = native_arguments(args, kwargs) + return rust_bridge.aocr(arguments) if arguments.get("aocr") is True else rust_bridge.ocr(arguments) + + async def private_aocr(*args, **kwargs): + if not configuration.rust_enabled(): + return await python_aocr(*args, **kwargs) + return await rust_bridge.aocr(native_arguments(args, kwargs)) + + monkeypatch.setattr(litellm, "ocr", private_ocr) + monkeypatch.setattr(litellm, "aocr", private_aocr) return native +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["ocr", "aocr", "ocr-async"]) +async def test_private_fixture_selects_native_only_when_enabled( + request, monkeypatch, injected_native, document, response, route +): + native = ModuleType("litellm.rust_bridge._native") + native.ocr = rust_bridge.ocr + native.aocr = rust_bridge.aocr + monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: native) + python_ocr, python_aocr = litellm.ocr, litellm.aocr + asynchronous = route != "ocr" + legacy = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(ocr_main, "_legacy_aocr" if route == "aocr" else "_legacy_ocr", legacy) + assert request.getfixturevalue("native_ocr") is native + assert ocr_main.ocr is python_ocr and ocr_main.aocr is python_aocr + arguments = {"metadata": {"opaque": object()}, **({"aocr": True} if route == "ocr-async" else {})} + for enabled in (True, False, True): + litellm.rust(enabled) + function = litellm.aocr if route == "aocr" else litellm.ocr + result = function(MODEL, document, **arguments) + assert (await result if asynchronous else result) is response + legacy.assert_called_once_with(MODEL, document, None, None, None, None, None, **arguments) + selected, unused = injected_native[int(asynchronous)], injected_native[not asynchronous] + assert selected.call_count == 2 + unused.assert_not_called() + assert selected.call_args.args[0] == { + "model": MODEL, + "document": document, + "api_key": None, + "api_base": None, + "timeout": None, + "custom_llm_provider": None, + "extra_headers": None, + **arguments, + } + assert selected.call_args.args[0]["document"] is document + assert selected.call_args.args[0]["metadata"] is arguments["metadata"] + + class WireRecorder: def __init__(self): self.requests = [] @@ -574,7 +625,7 @@ async def test_native_mistral_wire_response_and_callback_identity( @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("failure", [False, True], ids=["success", "failure"]) @pytest.mark.parametrize("callback_source", ["global", "per-call", "terminal-list"]) -async def test_public_native_callback_lifecycle( +async def test_private_native_callback_lifecycle( native_ocr, wire_recorder, monkeypatch, no_python_ocr, document, asynchronous, failure, callback_source ): from litellm import utils @@ -722,17 +773,59 @@ async def test_public_native_callback_lifecycle( @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True], ids=["success", "failure"]) -async def test_public_native_callable_terminal_callback( - native_ocr, wire_recorder, monkeypatch, no_python_ocr, document, failure +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync-request", "async-request"]) +@pytest.mark.parametrize("callback_kind", ["sync", "async", "async-object"]) +@pytest.mark.parametrize("callback_source", ["terminal-list", "global-list", "callbacks", "per-call", "overlap"]) +async def test_private_native_callable_terminal_callback( + native_ocr, + wire_recorder, + monkeypatch, + no_python_ocr, + document, + failure, + asynchronous, + callback_kind, + callback_source, ): + from litellm import utils + from litellm.litellm_core_utils import litellm_logging, logging_worker + litellm.rust(True) wire_recorder.status = 429 if failure else 200 calls = [] finished = threading.Event() + executor = ThreadPoolExecutor(max_workers=1) + worker = logging_worker.LoggingWorker(timeout=5, concurrency=1) + monkeypatch.setattr(utils, "executor", executor) + monkeypatch.setattr(litellm_logging, "executor", executor) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + monkeypatch.setattr(litellm_logging, "customLogger", None) + monkeypatch.setattr(utils, "callback_list", []) + monkeypatch.setattr(utils, "function_setup", Mock(side_effect=AssertionError("native OCR entered function_setup"))) - def callback(kwargs, response_obj, start_time, end_time): - calls.append((kwargs, response_obj, start_time, end_time)) - finished.set() + def record(kwargs, *terminal): + if terminal: + assert kwargs["log_event_type"] == "post_api_call" + calls.append((kwargs, *terminal)) + finished.set() + + async def async_callback(kwargs, *terminal): + await asyncio.sleep(0) + record(kwargs, *terminal) + + class AsyncCallable: + async def __call__(self, kwargs, *terminal): + await async_callback(kwargs, *terminal) + + callback = record if callback_kind == "sync" else async_callback if callback_kind == "async" else AsyncCallable() + callback_list = [callback, callback] + terminal_name = "failure_callback" if failure else "success_callback" + if callback_source in ("global-list", "overlap"): + monkeypatch.setattr(litellm, terminal_name, [callback]) + if callback_source == "overlap": + monkeypatch.setattr(litellm, f"_async_{terminal_name}", [callback]) + if callback_source == "callbacks": + monkeypatch.setattr(litellm, "callbacks", [callback]) arguments = { "model": MODEL, @@ -741,27 +834,91 @@ async def test_public_native_callable_terminal_callback( "api_base": wire_recorder.api_base, "timeout": 5, "num_retries": 0, - "failure_callback" if failure else "success_callback": [callback], + **({terminal_name: callback_list} if callback_source in ("terminal-list", "overlap") else {}), + **({"callbacks": callback_list} if callback_source in ("per-call", "overlap") else {}), } - if failure: - with pytest.raises(litellm.RateLimitError) as caught: - litellm.ocr(**arguments) - expected_response = None - else: - expected_response = litellm.ocr(**arguments) + try: + if failure: + with pytest.raises(litellm.RateLimitError) as caught: + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + expected_response = None + else: + expected_response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - assert await asyncio.to_thread(finished.wait, 5), "callable callback was not delivered" + if asynchronous or callback_kind == "sync": + assert await asyncio.to_thread(finished.wait, 5), "callable callback was not delivered" + finally: + try: + await asyncio.wait_for(worker.flush(), 5) + await asyncio.wait_for(asyncio.wrap_future(executor.submit(lambda: None)), 5) + finally: + await asyncio.wait_for(worker.stop(), 5) + atexit.unregister(worker._flush_on_exit) + executor.shutdown(wait=True, cancel_futures=True) + + assert callback_list == [callback, callback] + assert callback not in utils.callback_list + if callback_source in ("terminal-list", "overlap"): + assert arguments[terminal_name] is callback_list + if callback_source in ("per-call", "overlap"): + assert arguments["callbacks"] is callback_list + assert callback not in litellm.callbacks + assert callback not in litellm.input_callback + if not asynchronous and callback_kind != "sync": + assert calls == [] + return assert len(calls) == 1 details, callback_response, start_time, end_time = calls[0] assert details["model"] == "mistral-ocr-latest" assert details["litellm_call_id"] - assert details["log_event_type"] == "post_api_call" assert callback_response is expected_response assert start_time <= end_time if failure: assert details["exception"] is caught.value +@pytest.mark.asyncio +async def test_native_named_callback_initialization_preserves_aliases(monkeypatch, document): + from datetime import datetime + + from litellm.litellm_core_utils import litellm_logging + + events = [] + + class NamedLogger(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + events.append("input") + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + events.append("success") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + events.append("failure") + + callback = NamedLogger() + monkeypatch.setattr(litellm_logging, "_init_custom_logger_compatible_class", Mock(return_value=callback)) + monkeypatch.setattr(litellm, "input_callback", [callback]) + monkeypatch.setattr(litellm, "_async_success_callback", [callback]) + monkeypatch.setattr(litellm, "_async_failure_callback", [callback]) + callbacks = ["lago", callback, "lago"] + metadata = {"opaque": object()} + arguments = {"model": MODEL, "document": document, "callbacks": callbacks, "metadata": metadata} + logger = rust_bridge.initialize_logging(arguments, True) + assert arguments["document"] is document + assert arguments["metadata"] is metadata + assert arguments["callbacks"] is callbacks + assert callbacks == ["lago", callback, "lago"] + assert rust_bridge.initialize_logging(arguments, True) is logger + logger.pre_call(input="document", api_key="sk-test") + result = OCRResponse.model_validate(RESPONSE_DATA) + await logger.async_success_handler(result, datetime.now(), datetime.now()) + await logger.async_failure_handler(ValueError("test"), "test", datetime.now(), datetime.now()) + assert events == ["input", "success", "failure"] + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("public", [False, True], ids=["unwrapped", "public"]) @@ -773,9 +930,24 @@ async def test_public_native_callable_terminal_callback( ("azure_ai/mistral-ocr-latest", None, {}, "HTTP document URL to data URI conversion"), ("vertex_ai/mistral-ocr-latest", None, {}, "HTTP document URL to data URI conversion"), ("mistral-ocr-latest", "vertex_ai", {}, "HTTP document URL to data URI conversion"), - ("azure_ai/mistral-ocr-latest", None, {"api_key": None}, "Azure OCR credential acquisition"), - ("vertex_ai/mistral-ocr-latest", None, {"api_key": None}, "Vertex OCR credential acquisition"), - ("vertex_ai/deepseek-ocr-maas", None, {"api_key": None}, "Vertex OCR credential acquisition"), + ( + "azure_ai/mistral-ocr-latest", + None, + {"api_key": None, "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}, + "OCR credential acquisition", + ), + ( + "vertex_ai/mistral-ocr-latest", + None, + {"api_key": None, "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}, + "OCR credential acquisition", + ), + ( + "vertex_ai/deepseek-ocr-maas", + None, + {"api_key": None, "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}, + "OCR credential acquisition", + ), ("azure_ai/cohere/parse-v5.0", None, {}, "Cohere OCR request transformation"), ("cohere/parse-v5.0", None, {}, "OCR provider"), ("vertex_ai/deepseek-ocr-maas", None, {"stream": True}, "OCR streaming response handling"), @@ -825,7 +997,7 @@ async def test_native_unsupported_requests_never_prepare_or_send( @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("auth", ["key", "header", "environment"]) @pytest.mark.parametrize("provider", ["azure_ai", "vertex_ai", "deepseek"]) -async def test_public_native_cloud_wire_and_shallow_boundaries( +async def test_private_native_cloud_wire_and_shallow_boundaries( native_ocr, wire_recorder, monkeypatch, no_python_ocr, asynchronous, auth, provider ): monkeypatch.setenv("LITELLM_RUST", "1") @@ -917,7 +1089,7 @@ async def test_public_native_cloud_wire_and_shallow_boundaries( @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_public_native_azure_supplied_entra_token( +async def test_private_native_azure_supplied_entra_token( native_ocr, wire_recorder, monkeypatch, no_python_ocr, asynchronous ): monkeypatch.setenv("LITELLM_RUST", "1") @@ -991,13 +1163,393 @@ def test_native_sync_callback_reentry_without_event_loop( assert len(wire_recorder.requests) == 2 +@pytest.fixture +async def isolated_ocr_logging_worker(monkeypatch): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker(timeout=5, concurrency=1) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + yield worker + finally: + try: + await asyncio.wait_for(worker.flush(), 5) + finally: + try: + await asyncio.wait_for(worker.stop(), 5) + finally: + atexit.unregister(worker._flush_on_exit) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +@pytest.mark.parametrize("outcome", ["mutation", "replacement", "failure"]) +async def test_public_deployment_callback_parity( + native_ocr, wire_recorder, monkeypatch, document, enabled, outcome, isolated_ocr_logging_worker +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + events, captured, responses, snapshots, terminals = [], [], [], [], [] + metadata = {"deployment_state": {"phase": "caller"}} + state = metadata["deployment_state"] + replacement = OCRResponse.model_validate({**RESPONSE_DATA, "document_annotation": {"reviewed": True}}) + decoy_logger = object() + + class DeploymentLogger(CustomLogger): + def __init__(self, index): + super().__init__() + self.index = index + + async def async_pre_call_deployment_hook(self, kwargs, call_type): + assert kwargs["metadata"] is metadata + if self.index == 0: + assert kwargs["litellm_logging_obj"] is logger + captured.append(kwargs["litellm_logging_obj"]) + state["phase"] = "first" + return {**kwargs, "pages": [2], "litellm_logging_obj": decoy_logger} + assert kwargs["litellm_logging_obj"] is decoy_logger + assert state["phase"] == "first" and kwargs["pages"] == [2] + kwargs["pages"].append(3) + state["phase"] = "second" + events.append(("deployment_pre", call_type.value)) + + def log_pre_api_call(self, model, messages, kwargs): + assert kwargs is captured[0].model_call_details + assert kwargs["litellm_params"]["metadata"]["deployment_state"] is state + assert state["phase"] == "second" + events.append(("pre_api", kwargs["additional_args"]["complete_input_dict"].get("pages"))) + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + assert request_data["metadata"] is metadata + assert request_data["litellm_logging_obj"] is captured[0] + responses.append(response) + if self.index == 0: + response.document_annotation = {"mutated": True} + state["phase"] = "success" + return replacement if outcome == "replacement" else None + assert state["phase"] == "success" + assert response is (replacement if outcome == "replacement" else responses[0]) + response.document_annotation["second"] = True + events.append(("deployment_success", call_type.value)) + + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, **kwargs): + assert request_data["metadata"] is metadata + assert request_data["litellm_logging_obj"] is captured[0] + snapshots.append(exception) + if self.index == 0: + assert exception.status_code == 429 + state["phase"] = "error" + captured[0].model_call_details["deployment_error_state"] = state + exception.status_code = 418 + raise RuntimeError("observer failure must not replace the provider error") + assert exception is snapshots[0] and exception.status_code == 418 + assert captured[0].model_call_details["deployment_error_state"] is state + assert state["phase"] == "error" + events.append(("deployment_failure", call_type.value)) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + terminals.append((kwargs, response_obj)) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + terminals.append((kwargs, kwargs["exception"])) + + callbacks = [DeploymentLogger(0), DeploymentLogger(1)] + monkeypatch.setattr(litellm, "callbacks", callbacks) + logger = Logging( + model=MODEL, + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id=f"deployment-{enabled}-{outcome}", + function_id="", + dynamic_input_callbacks=callbacks, + dynamic_async_success_callbacks=callbacks, + dynamic_async_failure_callbacks=callbacks, + ) + wire_recorder.status = 429 if outcome == "failure" else 200 + litellm.rust(enabled) + arguments = dict( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + metadata=metadata, + litellm_logging_obj=logger, + ) + if outcome == "failure": + with pytest.raises(litellm.RateLimitError) as caught: + await litellm.aocr(**arguments) + result = caught.value + assert result.status_code == 429 + assert len(snapshots) == 2 and all(snapshot is not result for snapshot in snapshots) + else: + result = await litellm.aocr(**arguments) + assert len(responses) == 2 and result is responses[1] + assert (result is responses[0]) is (outcome == "mutation") + assert result.document_annotation == { + "reviewed" if outcome == "replacement" else "mutated": True, + "second": True, + } + await asyncio.sleep(0) + await asyncio.wait_for(isolated_ocr_logging_worker.flush(), 5) + assert len(captured) == 1 and captured[0] is logger + assert len(terminals) == 2 + for details, terminal in terminals: + assert details is captured[0].model_call_details and terminal is result + assert details["litellm_params"]["metadata"]["deployment_state"] is state + if outcome == "failure": + assert details["deployment_error_state"] is state + terminal_event = "deployment_failure" if outcome == "failure" else "deployment_success" + assert events == [("deployment_pre", "aocr"), ("pre_api", [2, 3]), ("pre_api", [2, 3]), (terminal_event, "aocr")] + assert len(wire_recorder.requests) == 1 and wire_recorder.requests[0]["body"]["pages"] == [2, 3] + + +@pytest.mark.asyncio +async def test_public_deferred_success_callback_parity( + native_ocr, wire_recorder, monkeypatch, document, isolated_ocr_logging_worker +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + worker = isolated_ocr_logging_worker + observations = [] + + class DeferredLogger(CustomLogger): + def __init__(self): + self.calls = [] + self.delivered = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.calls.append(response_obj.document_annotation) + self.delivered.set() + + try: + for enabled in (False, True): + callback = DeferredLogger() + logger = Logging( + model=MODEL, + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id=f"deferred-{enabled}", + function_id="", + dynamic_async_success_callbacks=[callback], + ) + logger._defer_async_logging = True + litellm.rust(enabled) + result = await litellm.aocr( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + litellm_logging_obj=logger, + ) + enqueue = getattr(logger, "_enqueue_deferred_logging", None) + assert callable(enqueue) + assert not callback.calls + result.document_annotation = {"reviewed": True} + enqueue() + logger._enqueue_deferred_logging = None + await asyncio.wait_for(callback.delivered.wait(), 5) + await asyncio.wait_for(worker.flush(), 5) + observations.append((callable(enqueue), tuple(callback.calls))) + finally: + await asyncio.wait_for(worker.flush(), 5) + assert len(wire_recorder.requests) == 2 + assert observations[0] == (True, ({"reviewed": True},)) + assert len(observations[1][1]) == 1 + assert observations[1] == observations[0] + + +def test_public_cold_callable_input_callback_parity(native_ocr, wire_recorder, monkeypatch, document): + from litellm import utils + from litellm.litellm_core_utils import litellm_logging + + observations = [] + for enabled in (False, True): + calls = [] + + def callback(kwargs, calls=calls): + calls.append(kwargs["log_event_type"]) + + monkeypatch.setattr(litellm, "input_callback", [callback]) + monkeypatch.setattr(litellm_logging, "customLogger", None) + monkeypatch.setattr(utils, "callback_list", []) + litellm.rust(enabled) + result = litellm.ocr( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + ) + assert result.pages[0].markdown == "proof" + observations.append(tuple(calls)) + assert len(wire_recorder.requests) == 2 + assert observations[0] == ("pre_api_call",) + assert observations[1] == observations[0] + + class Opaque: pass +@pytest.mark.asyncio +async def test_native_deferred_callback_retains_shared_arguments( + native_ocr, wire_recorder, document, isolated_ocr_logging_worker +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + class Arguments(dict): + pass + + entered, release = asyncio.Event(), asyncio.Event() + observed = [] + + class DeferredLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + entered.set() + await release.wait() + shared = arguments_ref() + observed.append( + ( + shared is not None and shared["opaque"] is opaque_ref(), + shared is not None and shared["unknown_option"]["nested"] is opaque_ref(), + shared is not None and shared["document"] is document, + response_obj is result, + ) + ) + + logger = Logging( + model=MODEL, + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id="retained-deferred", + function_id="", + dynamic_async_success_callbacks=[DeferredLogger()], + ) + logger._defer_async_logging = True + opaque = Opaque() + arguments = Arguments( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + litellm_logging_obj=logger, + opaque=opaque, + unknown_option={"nested": opaque}, + ) + arguments_ref, opaque_ref = weakref.ref(arguments), weakref.ref(opaque) + try: + result = await native_ocr.aocr(arguments) + assert result.pages[0].markdown == "proof" and not entered.is_set() + logger._enqueue_deferred_logging() + logger._enqueue_deferred_logging = None + del arguments, opaque + await asyncio.wait_for(entered.wait(), 5) + gc.collect() + assert arguments_ref() is not None and opaque_ref() is not None + assert arguments_ref()["unknown_option"]["nested"] is opaque_ref() + assert not observed + shared = arguments_ref() + finally: + release.set() + await asyncio.wait_for(isolated_ocr_logging_worker.flush(), 5) + assert observed == [(True, True, True, True)] + assert shared["opaque"] is shared["unknown_option"]["nested"] is opaque_ref() + assert shared["document"] is document and shared["litellm_logging_obj"] is logger + del shared + await asyncio.sleep(0) + gc.collect() + assert arguments_ref() is None and opaque_ref() is None + assert len(wire_recorder.requests) == 1 + assert "opaque" not in wire_recorder.requests[0]["body"] + assert "unknown_option" not in wire_recorder.requests[0]["body"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +@pytest.mark.parametrize("outcome", ["success", "failure", "cancel"]) +async def test_correlation_context_restored_in_calling_task( + native_ocr, wire_recorder, monkeypatch, document, enabled, outcome, isolated_ocr_logging_worker +): + from litellm._logging import session_id_var, trace_id_var + + during, restored = [], [] + + class CorrelationLogger(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + during.append((trace_id_var.get(), session_id_var.get(), asyncio.current_task())) + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + monkeypatch.setattr(litellm, "callbacks", [CorrelationLogger()]) + litellm.rust(enabled) + wire_recorder.status = 429 if outcome == "failure" else 200 + wire_recorder.release.clear() + + async def call(): + trace_token = trace_id_var.set("outer-trace") + session_token = session_id_var.set("outer-session") + try: + return await litellm.aocr( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + litellm_trace_id="request-trace", + litellm_session_id="request-session", + ) + finally: + restored.append((trace_id_var.get(), session_id_var.get(), asyncio.current_task())) + trace_id_var.reset(trace_token) + session_id_var.reset(session_token) + + task = asyncio.create_task(call()) + try: + assert await asyncio.to_thread(wire_recorder.received.wait, 5), "POST never reached server" + assert during == [("request-trace", "request-session", task)] + if outcome == "cancel": + task.cancel() + else: + wire_recorder.release.set() + if outcome == "success": + assert (await asyncio.wait_for(task, 5)).pages[0].markdown == "proof" + else: + with pytest.raises(asyncio.CancelledError if outcome == "cancel" else litellm.RateLimitError): + await asyncio.wait_for(task, 5) + assert restored == [("outer-trace", "outer-session", task)] + finally: + wire_recorder.release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert await asyncio.to_thread(wire_recorder.finished.wait, 5) + await asyncio.wait_for(isolated_ocr_logging_worker.flush(), 5) + + @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["success", "error", "cancel"]) -async def test_native_retains_opaque_arguments_until_terminal_cleanup(native_ocr, wire_recorder, document, outcome): +async def test_native_retains_opaque_arguments_until_terminal_cleanup( + native_ocr, wire_recorder, document, outcome, isolated_ocr_logging_worker +): wire_recorder.release.clear() wire_recorder.status = 429 if outcome == "error" else 200 opaque = Opaque() @@ -1046,10 +1598,8 @@ async def test_native_retains_opaque_arguments_until_terminal_cleanup(native_ocr await asyncio.gather(task, return_exceptions=True) assert await asyncio.to_thread(wire_recorder.finished.wait, 5) del task - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - if outcome == "success": - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5) + await asyncio.wait_for(isolated_ocr_logging_worker.flush(), 5) logger.reset_mock() await asyncio.sleep(0) gc.collect() @@ -1102,3 +1652,158 @@ async def test_public_cancellation_does_not_emit_failure_callbacks( await asyncio.sleep(0) assert calls == ["pre_call"] assert len(wire_recorder.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +async def test_public_concurrent_callback_isolation_and_cleanup(request, monkeypatch, enabled): + from contextlib import ExitStack + + from litellm import utils + from litellm.litellm_core_utils import litellm_logging, logging_worker + + if enabled: + request.getfixturevalue("native_ocr") + request.getfixturevalue("no_python_ocr") + litellm.rust(enabled) + context = contextvars.ContextVar("ocr-stress-context", default="parent") + pre_calls, terminals = [], [] + entered, release = {}, {} + worker = logging_worker.LoggingWorker(timeout=5, concurrency=2) + executor = ThreadPoolExecutor(max_workers=1) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + monkeypatch.setattr(utils, "executor", executor) + monkeypatch.setattr(litellm_logging, "executor", executor) + + class ConcurrentLogger(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + tag = kwargs["litellm_call_id"] + state = {"owner": tag, "phase": "pre"} + pre_calls.append((tag, kwargs, state, context.get())) + kwargs["ocr_stress_state"] = state + kwargs["additional_args"]["headers"]["X-Callback"] = tag + context.set(f"{tag}:pre") + + def record(self, event, kwargs, response_obj): + tag = kwargs["litellm_call_id"] + state = kwargs.get("ocr_stress_state", {}) + terminals.append( + (tag, event, kwargs, state, dict(state), context.get(), response_obj, kwargs.get("exception")) + ) + state["phase"] = event + context.set(f"{tag}:{event}") + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.record("sync_success", kwargs, response_obj) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.record("sync_failure", kwargs, response_obj) + + async def terminal(self, event, kwargs, response_obj): + tag = kwargs["litellm_call_id"] + entered[tag].set() + await asyncio.wait_for(release[tag].wait(), 5) + self.record(event, kwargs, response_obj) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self.terminal("async_success", kwargs, response_obj) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await self.terminal("async_failure", kwargs, response_obj) + + monkeypatch.setattr(litellm, "callbacks", [ConcurrentLogger()]) + + async def call(tag, recorder): + context.set(tag) + return await litellm.aocr( + model=MODEL, + document={"type": "document_url", "document_url": f"https://example.invalid/{tag}.pdf"}, + api_key="sk-test", + api_base=recorder.api_base, + litellm_call_id=tag, + timeout=5, + num_retries=0, + ) + + async def batch(prefix, outcomes): + with ExitStack() as stack: + recorders = {f"{prefix}-{outcome}": WireRecorder() for outcome in outcomes} + for tag, recorder in recorders.items(): + stack.callback(recorder.stop) + recorder.release.clear() + recorder.status = 429 if tag.endswith("error") else 200 + entered[tag], release[tag] = asyncio.Event(), asyncio.Event() + tasks = {tag: asyncio.create_task(call(tag, recorder)) for tag, recorder in recorders.items()} + try: + assert all(await asyncio.gather(*(asyncio.to_thread(r.received.wait, 5) for r in recorders.values()))) + assert all(not task.done() for task in tasks.values()) + for tag, task in tasks.items(): + if tag.endswith("cancel"): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 5) + else: + recorders[tag].release.set() + active = tuple(tag for tag in tasks if not tag.endswith("cancel")) + await asyncio.wait_for(asyncio.gather(*(entered[tag].wait() for tag in active)), 5) + for tag in reversed(active): + release[tag].set() + results = await asyncio.wait_for(asyncio.gather(*tasks.values(), return_exceptions=True), 5) + await asyncio.wait_for(worker.flush(), 5) + await asyncio.wait_for(asyncio.wrap_future(executor.submit(lambda: None)), 5) + for (tag, recorder), result in zip(recorders.items(), results): + matching_pre = [entry for entry in pre_calls if entry[0] == tag] + assert len(matching_pre) == len(recorder.requests) == 1 + _, details, state, pre_context = matching_pre[0] + assert pre_context == tag + sent = recorder.requests[0] + assert sent["path"] == "/v1/ocr" + assert sent["headers"]["x-callback"] == tag + assert sent["body"]["document"]["document_url"] == f"https://example.invalid/{tag}.pdf" + expected = ( + [] + if tag.endswith("cancel") + else ["sync_failure", "async_failure"] + if tag.endswith("error") + else ["async_success"] + ) + matching_terminal = [entry for entry in terminals if entry[0] == tag] + assert [entry[1] for entry in matching_terminal] == expected + if tag.endswith("cancel"): + assert isinstance(result, asyncio.CancelledError) + elif tag.endswith("error"): + assert isinstance(result, litellm.RateLimitError) and result.status_code == 429 + else: + assert isinstance(result, OCRResponse) and result.pages[0].markdown == "proof" + for _, event, kwargs, shared, snapshot, terminal_context, response, error in matching_terminal: + phase = "sync_failure" if event == "async_failure" else "pre" + assert kwargs is details and shared is state + assert snapshot == {"owner": tag, "phase": phase} + assert terminal_context == f"{tag}:{phase}" + assert response is (None if tag.endswith("error") else result) + assert error is (result if tag.endswith("error") else None) + assert context.get() == "parent" + finally: + for tag, recorder in recorders.items(): + recorder.release.set() + release[tag].set() + if not tasks[tag].done(): + tasks[tag].cancel() + await asyncio.wait_for(asyncio.gather(*tasks.values(), return_exceptions=True), 5) + assert all(await asyncio.gather(*(asyncio.to_thread(r.finished.wait, 5) for r in recorders.values()))) + + try: + for index in range(2): + await batch(str(index), ("success", "error", "cancel")) + await batch("recovery", ("success",)) + assert len(pre_calls) == 7 and len(terminals) == 7 + assert len({id(entry[1]) for entry in pre_calls}) == len({id(entry[2]) for entry in pre_calls}) == 7 + finally: + try: + await asyncio.wait_for(worker.flush(), 5) + finally: + try: + await asyncio.wait_for(worker.stop(), 5) + finally: + atexit.unregister(worker._flush_on_exit) + await asyncio.wait_for(asyncio.to_thread(executor.shutdown, wait=True, cancel_futures=True), 5) diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 7427eba5fa5..01deaeca2fb 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -10,6 +10,8 @@ import sys import tempfile import threading import zipfile +from collections.abc import Awaitable +from contextvars import ContextVar from dataclasses import dataclass from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -200,10 +202,35 @@ def initialize_ocr_logging(arguments: dict[str, object], asynchronous: bool) -> def invoke_ocr_terminal( action: str, roots: object, logger: OCRLogging, value: object, start: object, end: object -) -> None: +) -> Awaitable[None] | None: assert action in {"sync_success", "async_success", "sync_success_if_needed", "sync_failure", "async_failure"} assert isinstance(roots, tuple) and roots[0] is logger.arguments assert logger.calls == ("update", "pre") + if action == "async_failure": + return observe_ocr_failure({}, value, "aocr") + return None + + +async def pre_ocr_deployment(arguments: dict[str, object], call_type: str) -> dict[str, object]: + assert call_type == "aocr" + return arguments + + +async def post_ocr_deployment(arguments: dict[str, object], response: object, call_type: str) -> object: + assert call_type == "aocr" + return response + + +async def observe_ocr_failure(arguments: dict[str, object], error: object, call_type: str) -> None: + assert call_type == "aocr" + + +def restore_ocr_context(logger: object) -> None: + pass + + +class WheelCallTypes: + aocr = "aocr" def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: @@ -369,11 +396,20 @@ def exercise_routes(native_path: Path, api_base: str) -> object: ocr_bridge: Final = ModuleType("litellm.rust_bridge.ocr") ocr_bridge.initialize_logging = initialize_ocr_logging ocr_bridge.invoke_terminal = invoke_ocr_terminal + utils: Final = ModuleType("litellm.utils") + utils.is_internal_call = ContextVar("wheel_internal_call", default=False) + utils.async_pre_call_deployment_hook = pre_ocr_deployment + utils.async_post_call_success_deployment_hook = post_ocr_deployment + utils.async_post_call_failure_deployment_hook = observe_ocr_failure + utils._restore_correlation_context_if_supported = restore_ocr_context + types_utils: Final = ModuleType("litellm.types.utils") + types_utils.CallTypes = WheelCallTypes packages: Final = { name: ModuleType(name) for name in ( "litellm", "litellm.rust_bridge", + "litellm.types", "litellm.llms", "litellm.llms.base_llm", "litellm.llms.base_llm.ocr", @@ -381,7 +417,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object: } with patch.dict( sys.modules, - packages | {module.__name__: module for module in (transformation, exceptions, httpx, ocr_bridge)}, + packages + | {module.__name__: module for module in (transformation, exceptions, httpx, ocr_bridge, utils, types_utils)}, ): exercise_sync(native, api_base) asyncio.run(exercise_async(native, api_base))