diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c2c1e454e6a..55bb5318cfa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1482,12 +1482,15 @@ dependencies = [ name = "litellm-python-interop" version = "0.1.0" dependencies = [ + "futures-util", "pyo3", + "pyo3-async-runtimes", "pythonize", "rstest", "serde", "serde_json", "serial_test", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 9fef59a277d..28d94bf16e8 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -13,7 +13,7 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | | litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-interop | Domain-neutral PyO3 foundation: typed Python/Serde conversion, retained callbacks, and sync/async Python↔Tokio execution. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 637c156e192..cd84c05a510 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -2,6 +2,7 @@ use std::future::Future; use std::time::{Instant, SystemTime, UNIX_EPOCH}; use crate::Error; +use crate::lifecycle::ActionResult; pub mod types; @@ -10,6 +11,39 @@ pub use types::{ CallLifecycleTiming, }; +pub trait RequestPolicy: Send + Sync { + type PreCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a; + + type DuringCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a; + + fn async_pre_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::PreCallFuture<'a>; + + fn async_during_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::DuringCallFuture<'a>; +} + +pub trait TerminalDispatcher: Send + Sync { + fn dispatch<'a>( + &'a self, + terminal: &'a crate::lifecycle::TerminalRecord, + ) -> crate::integrations::custom_logger::LogFuture<'a>; +} + pub trait CallLifecycleHooks: Send + Sync { type PreCallFuture<'a>: Future> + Send + 'a where @@ -97,6 +131,21 @@ impl<'a> CallLifecycle<'a> { self.run(context, request, hooks, provider_call).await } + pub async fn run_result( + &self, + context: CallLifecycleContext, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> Result + where + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + self.run(context, request, hooks, provider_call).await + } + pub async fn run( &self, context: CallLifecycleContext, diff --git a/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs index 792717dacfc..6430c7cbbff 100644 --- a/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs +++ b/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs @@ -96,6 +96,36 @@ impl CustomLoggerRunner { } } +impl crate::call_lifecycle::TerminalDispatcher for CustomLoggerRunner { + fn dispatch<'a>(&'a self, terminal: &'a crate::lifecycle::TerminalRecord) -> LogFuture<'a> { + Box::pin(async move { + let details = ModelCallDetails::from(terminal); + let response = CallbackValue::new( + match &terminal.projection { + crate::lifecycle::RouteProjection::Ocr { .. } => "ocr", + crate::lifecycle::RouteProjection::Messages { .. } => "messages", + crate::lifecycle::RouteProjection::ChatCompletions { .. } => "chat_completion", + crate::lifecycle::RouteProjection::Audio { .. } => "audio", + crate::lifecycle::RouteProjection::Realtime { .. } => "realtime", + crate::lifecycle::RouteProjection::ResponsesWs { .. } => "responses_websocket", + }, + terminal.projection.value().clone(), + ); + match terminal.classification { + crate::lifecycle::TerminalClassification::Success => { + self.async_log_success_event(&details, &response, terminal.timing) + .await; + } + crate::lifecycle::TerminalClassification::Failure { .. } => { + self.async_log_failure_event(&details, Some(&response), terminal.timing) + .await; + } + } + Ok(()) + }) + } +} + #[cfg(test)] mod tests { use super::*; @@ -303,6 +333,38 @@ mod tests { assert_eq!(report, CallbackDispatchReport::default()); } + #[tokio::test] + async fn terminal_dispatcher_fans_out_shared_record() { + use crate::call_lifecycle::TerminalDispatcher; + use crate::integrations::types::Usage; + use crate::lifecycle::terminal::CostInputs; + use crate::lifecycle::{RouteProjection, TerminalClassification, TerminalRecord}; + + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let terminal = TerminalRecord { + call_id: "call-ocr".to_string(), + trace_id: Some("req-ocr".to_string()), + attempt: 1, + call_type: "ocr".to_string(), + model: "mistral-ocr-latest".to_string(), + provider: "mistral".to_string(), + timing: CallbackTiming::new(10.0, 11.5), + usage: Usage::default(), + cost_inputs: CostInputs::default(), + classification: TerminalClassification::Success, + projection: RouteProjection::Ocr { + value: json!({"pages": []}), + }, + }; + + runner.dispatch(&terminal).await.expect("dispatch succeeds"); + + assert_eq!(logger.events().len(), 1); + assert_eq!(logger.events()[0].hook, "async_log_success_event"); + assert_eq!(logger.events()[0].response_object.as_deref(), Some("ocr")); + } + #[test] fn with_standard_logging_payload_keeps_top_level_fields_in_sync() { let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion) diff --git a/litellm-rust/crates/core/src/integrations/custom_logger/types.rs b/litellm-rust/crates/core/src/integrations/custom_logger/types.rs index ba7d67bd46e..e68babc7821 100644 --- a/litellm-rust/crates/core/src/integrations/custom_logger/types.rs +++ b/litellm-rust/crates/core/src/integrations/custom_logger/types.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; +use serde::Serialize; use serde_json::Value; use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; @@ -56,7 +57,7 @@ impl std::fmt::Display for CallType { } } -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Copy, Debug, PartialEq, Serialize)] pub struct CallbackTiming { pub start_time: f64, pub end_time: f64, diff --git a/litellm-rust/crates/core/src/integrations/types.rs b/litellm-rust/crates/core/src/integrations/types.rs index 34dce93d8e0..a523e14fd24 100644 --- a/litellm-rust/crates/core/src/integrations/types.rs +++ b/litellm-rust/crates/core/src/integrations/types.rs @@ -13,7 +13,7 @@ use serde_json::Value; use std::collections::HashMap; /// Cumulative token usage for a realtime session. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] pub struct Usage { pub prompt_tokens: u64, pub completion_tokens: u64, @@ -65,7 +65,7 @@ pub struct StandardLoggingPayload { /// Cost-attribution keys. The replayer maps these into litellm_params.metadata, /// which the spend-logs builder reads to set user / team_id / organization_id. -#[derive(Clone, Debug, Serialize, Default)] +#[derive(Clone, Debug, Serialize, Default, PartialEq)] pub struct StandardLoggingMetadata { pub user_api_key_hash: Option, // -> SpendLogs.api_key pub user_api_key_user_id: Option, // -> SpendLogs.user diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 25f4fce6f43..2d9ebb836b2 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -6,6 +6,7 @@ pub mod constants; pub mod error; pub mod http_utils; pub mod integrations; +pub mod lifecycle; pub mod messages; #[cfg(any(feature = "observability", test))] pub mod observability; diff --git a/litellm-rust/crates/core/src/lifecycle/action.rs b/litellm-rust/crates/core/src/lifecycle/action.rs new file mode 100644 index 00000000000..750ded6198f --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/action.rs @@ -0,0 +1,26 @@ +use serde::Serialize; + +use super::{ActionKind, Delivery, FailurePolicy}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum ResultPolicy { + Continue, + Replace, + Reject, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum Owner { + Core, + Route, + Host, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub struct ActionBinding { + pub kind: ActionKind, + pub delivery: Delivery, + pub on_result: ResultPolicy, + pub on_error: FailurePolicy, + pub owner: Owner, +} diff --git a/litellm-rust/crates/core/src/lifecycle/executed.rs b/litellm-rust/crates/core/src/lifecycle/executed.rs new file mode 100644 index 00000000000..072e0177aa4 --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/executed.rs @@ -0,0 +1,64 @@ +use super::TerminalRecord; + +#[derive(Clone, Debug)] +pub enum ExecutedCall { + Success { + response: R, + terminal: TerminalRecord, + }, + Failure { + error: E, + terminal: TerminalRecord, + }, +} + +impl ExecutedCall { + pub fn into_result(self) -> Result { + match self { + Self::Success { response, .. } => Ok(response), + Self::Failure { error, .. } => Err(error), + } + } + + pub fn terminal(&self) -> &TerminalRecord { + match self { + Self::Success { terminal, .. } | Self::Failure { terminal, .. } => terminal, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::custom_logger::CallbackTiming; + use crate::integrations::types::Usage; + use crate::lifecycle::terminal::CostInputs; + use crate::lifecycle::{RouteProjection, TerminalClassification}; + use serde_json::json; + + #[test] + fn failure_retains_its_terminal_record() { + let call = ExecutedCall::<(), _>::Failure { + error: "provider failed", + terminal: TerminalRecord { + call_id: "call-1".to_string(), + trace_id: None, + attempt: 1, + call_type: "ocr".to_string(), + model: "model".to_string(), + provider: "provider".to_string(), + timing: CallbackTiming::new(1.0, 2.0), + usage: Usage::default(), + cost_inputs: CostInputs::default(), + classification: TerminalClassification::Failure { + kind: "ProviderError".to_string(), + message: "provider failed".to_string(), + }, + projection: RouteProjection::Ocr { value: json!({}) }, + }, + }; + + assert_eq!(call.terminal().call_id, "call-1"); + assert_eq!(call.into_result(), Err("provider failed")); + } +} diff --git a/litellm-rust/crates/core/src/lifecycle/machine.rs b/litellm-rust/crates/core/src/lifecycle/machine.rs new file mode 100644 index 00000000000..a9cbbb502f0 --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/machine.rs @@ -0,0 +1,68 @@ +use std::marker::PhantomData; + +use super::ActionBinding; + +pub trait LifecycleRoute: Sized { + type Admission; + type Options; + type Context; + type Operation: Copy; + type Observation; + type Outcome; + type Transition; + type Error; + type Decline; + type State; + + fn admit( + admission: &Self::Admission, + options: Self::Options, + ) -> Result, Self::Error>; + + fn operation(state: &Self::State) -> Self::Operation; + + fn advance( + state: &mut Self::State, + outcome: Self::Outcome, + observations: Self::Observation, + ) -> Result; + + fn actions_for(operation: Self::Operation, context: &Self::Context) + -> &'static [ActionBinding]; +} + +#[derive(Debug)] +pub struct Lifecycle { + pub(crate) state: Route::State, + route: PhantomData, +} + +impl Lifecycle { + pub fn admit( + admission: &Route::Admission, + options: Route::Options, + ) -> Result, Route::Error> { + Route::admit(admission, options).map(|admission| { + admission.map(|state| Self { + state, + route: PhantomData, + }) + }) + } + + pub fn operation(&self) -> Route::Operation { + Route::operation(&self.state) + } + + pub fn advance( + &mut self, + outcome: Route::Outcome, + observations: Route::Observation, + ) -> Result { + Route::advance(&mut self.state, outcome, observations) + } + + pub fn actions_for(&self, context: &Route::Context) -> &'static [ActionBinding] { + Route::actions_for(self.operation(), context) + } +} diff --git a/litellm-rust/crates/core/src/lifecycle/mod.rs b/litellm-rust/crates/core/src/lifecycle/mod.rs new file mode 100644 index 00000000000..9292999759f --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/mod.rs @@ -0,0 +1,12 @@ +pub mod action; +pub mod executed; +pub mod machine; +pub mod ocr; +pub mod terminal; +pub mod types; + +pub use action::{ActionBinding, Owner, ResultPolicy}; +pub use executed::ExecutedCall; +pub use machine::{Lifecycle, LifecycleRoute}; +pub use terminal::{RouteProjection, TerminalClassification, TerminalRecord}; +pub use types::{ActionKind, ActionResult, Delivery, ErrorDisposition, FailurePolicy, Outcome}; diff --git a/litellm-rust/crates/core/src/lifecycle/ocr.rs b/litellm-rust/crates/core/src/lifecycle/ocr.rs new file mode 100644 index 00000000000..d3eba181b9f --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/ocr.rs @@ -0,0 +1,264 @@ +use crate::Error; +use crate::ocr::{OcrRequest, prepare}; + +use super::{ + ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, LifecycleRoute, Outcome, + Owner, ResultPolicy, +}; + +#[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 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 struct Transition { + pub operation: Operation, + pub error: ErrorDisposition, +} + +#[derive(Debug)] +pub struct OcrState { + operation: Operation, + outcome: Outcome, + asynchronous: bool, + internal_call: bool, + identity: Identity, +} + +#[derive(Debug)] +pub struct OcrRoute; + +pub type Lifecycle = super::Lifecycle; + +impl Lifecycle { + pub fn new(admission: &OcrRequest, options: Options) -> Result, Error> { + >::admit(admission, options).map(|admission| match admission { + Ok(lifecycle) => NativeOutcome::Completed(lifecycle), + Err(decline) => NativeOutcome::Declined(decline), + }) + } + + pub fn identity(&self) -> &Identity { + &self.state.identity + } +} + +impl LifecycleRoute for OcrRoute { + type Admission = OcrRequest; + type Options = Options; + type Context = Observations; + type Operation = Operation; + type Observation = Observations; + type Outcome = Outcome; + type Transition = Transition; + type Error = Error; + type Decline = Decline; + type State = OcrState; + + fn admit( + admission: &Self::Admission, + options: Self::Options, + ) -> Result, Self::Error> { + match prepare::admission_capabilities(admission) { + Err(Error::Unsupported(reason)) => return Ok(Err(Decline(reason))), + Err(error) => return Err(error), + Ok(()) => {} + } + if options.credential_method == CredentialMethod::Acquisition { + return Ok(Err(Decline("OCR credential acquisition"))); + } + let generated_call_id = options.call_id.is_none(); + let call_id = options.call_id.unwrap_or_else(generate_call_id); + Ok(Ok(OcrState { + 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, + }, + })) + } + + fn operation(state: &Self::State) -> Self::Operation { + state.operation + } + + fn advance( + state: &mut Self::State, + outcome: Self::Outcome, + observations: Self::Observation, + ) -> Result { + use Operation::*; + + if matches!(state.operation, Complete(_)) { + return Err(Error::InvalidRequest( + "OCR lifecycle is already complete".into(), + )); + } + let failure = + if observations.logger_available && !(state.asynchronous && state.internal_call) { + SyncFailure + } else { + Restore + }; + let error = if outcome != Outcome::Success && state.operation != DeploymentFailure { + state.outcome = outcome; + ErrorDisposition::Replace + } else { + ErrorDisposition::Preserve + }; + state.operation = match (state.operation, outcome) { + (Restore, _) => Complete(state.outcome), + (DeploymentFailure, _) => failure, + (_, Outcome::Abort) => Restore, + (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, + (Prepare | Send, Outcome::Failure) if state.asynchronous => DeploymentFailure, + (_, Outcome::Failure) => failure, + (Setup, Outcome::Success) if state.asynchronous => DeploymentPre, + (Setup | DeploymentPre, Outcome::Success) => Prepare, + (Prepare, Outcome::Success) => Send, + (Send, Outcome::Success) if state.asynchronous => DeploymentSuccess, + (Send, Outcome::Success) => SyncSuccess, + (DeploymentSuccess, Outcome::Success) => { + if state.internal_call || observations.has_fallbacks { + SyncSuccessIfNeeded + } else { + AsyncSuccess + } + } + (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, + (SyncFailure, Outcome::Success) if state.asynchronous => AsyncFailure, + (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { + Restore + } + (Complete(_), _) => unreachable!(), + }; + Ok(Transition { + operation: state.operation, + error, + }) + } + + fn actions_for( + operation: Self::Operation, + _context: &Self::Context, + ) -> &'static [ActionBinding] { + match operation { + Operation::Prepare | Operation::Send => &PROVIDER_ACTION, + Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { + &FAILURE_ACTION + } + Operation::Restore => &RESTORE_ACTION, + Operation::Complete(_) => &[], + _ => &CALLBACK_ACTION, + } + } +} + +const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::ProviderCall, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Replace, + on_error: FailurePolicy::Propagate, + owner: Owner::Core, +}]; + +const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::TerminalSuccess, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::RecordAndContinue, + owner: Owner::Route, +}]; + +const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::TerminalFailure, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::PreserveOriginalFailure, + owner: Owner::Route, +}]; + +const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::Restore, + delivery: Delivery::InlineDirect, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::Propagate, + owner: Owner::Core, +}]; + +fn generate_call_id() -> String { + 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..] + ) +} diff --git a/litellm-rust/crates/core/src/lifecycle/terminal.rs b/litellm-rust/crates/core/src/lifecycle/terminal.rs new file mode 100644 index 00000000000..63c2869254c --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/terminal.rs @@ -0,0 +1,156 @@ +use serde::Serialize; +use serde_json::Value; + +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload, Usage}; + +#[derive(Clone, Debug, Default, PartialEq, Serialize)] +pub struct CostInputs { + pub response_cost: f64, + pub metadata: StandardLoggingMetadata, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub enum TerminalClassification { + Success, + Failure { kind: String, message: String }, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub enum RouteProjection { + Ocr { value: Value }, + Messages { value: Value }, + ChatCompletions { value: Value }, + Audio { value: Value }, + Realtime { value: Value }, + ResponsesWs { value: Value }, +} + +impl RouteProjection { + pub(crate) fn value(&self) -> &Value { + match self { + Self::Ocr { value } + | Self::Messages { value } + | Self::ChatCompletions { value } + | Self::Audio { value } + | Self::Realtime { value } + | Self::ResponsesWs { value } => value, + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct TerminalRecord { + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub call_type: String, + pub model: String, + pub provider: String, + pub timing: CallbackTiming, + pub usage: Usage, + pub cost_inputs: CostInputs, + pub classification: TerminalClassification, + pub projection: RouteProjection, +} + +impl From<&TerminalRecord> for StandardLoggingPayload { + fn from(record: &TerminalRecord) -> Self { + Self { + id: record + .trace_id + .clone() + .unwrap_or_else(|| record.call_id.clone()), + litellm_call_id: record.call_id.clone(), + call_type: record.call_type.clone(), + model: record.model.clone(), + custom_llm_provider: record.provider.clone(), + response_cost: record.cost_inputs.response_cost, + prompt_tokens: record.usage.prompt_tokens, + completion_tokens: record.usage.completion_tokens, + total_tokens: record.usage.total_tokens, + start_time: record.timing.start_time, + end_time: record.timing.end_time, + stream: matches!( + record.projection, + RouteProjection::Realtime { .. } | RouteProjection::ResponsesWs { .. } + ), + metadata: record.cost_inputs.metadata.clone(), + messages: Some(record.projection.value().clone()), + } + } +} + +impl From<&TerminalRecord> for ModelCallDetails { + fn from(record: &TerminalRecord) -> Self { + let details = Self::from_standard_logging_payload(record.into()); + match &record.classification { + TerminalClassification::Success => details, + TerminalClassification::Failure { kind, message } => { + details.with_failure_error(LoggingError { + kind: kind.clone(), + message: message.clone(), + }) + } + } + } +} + +impl TerminalRecord { + pub fn call_type(&self) -> CallType { + CallType::from(self.call_type.as_str()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn record() -> TerminalRecord { + TerminalRecord { + call_id: "call-1".to_string(), + trace_id: Some("trace-1".to_string()), + attempt: 2, + call_type: "ocr".to_string(), + model: "mistral-ocr-latest".to_string(), + provider: "mistral".to_string(), + timing: CallbackTiming::new(10.0, 11.5), + usage: Usage { + prompt_tokens: 3, + completion_tokens: 4, + total_tokens: 7, + }, + cost_inputs: CostInputs { + response_cost: 0.25, + metadata: StandardLoggingMetadata { + user_api_key_user_id: Some("user-1".to_string()), + ..Default::default() + }, + }, + classification: TerminalClassification::Success, + projection: RouteProjection::Ocr { + value: json!({"pages": [{"markdown": "ok"}]}), + }, + } + } + + #[test] + fn terminal_record_projects_existing_logging_payload() { + let record = record(); + let payload = StandardLoggingPayload::from(&record); + let details = ModelCallDetails::from(&record); + + assert_eq!(payload.id, "trace-1"); + assert_eq!(payload.litellm_call_id, "call-1"); + assert_eq!(payload.prompt_tokens, 3); + assert_eq!(payload.response_cost, 0.25); + assert_eq!(payload.messages, Some(record.projection.value().clone())); + assert_eq!( + details.standard_logging_payload.unwrap().litellm_call_id, + "call-1" + ); + } +} diff --git a/litellm-rust/crates/core/src/lifecycle/types.rs b/litellm-rust/crates/core/src/lifecycle/types.rs new file mode 100644 index 00000000000..0c089cbb377 --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/types.rs @@ -0,0 +1,56 @@ +use serde::Serialize; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum ActionKind { + RequestPolicy, + ProviderCall, + Deployment, + TerminalSuccess, + TerminalFailure, + Restore, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ActionResult { + Continue(T), + Replace(T), + Reject(E), +} + +impl ActionResult { + pub fn into_result(self) -> Result { + match self { + Self::Continue(value) | Self::Replace(value) => Ok(value), + Self::Reject(error) => Err(error), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum FailurePolicy { + Propagate, + RecordAndContinue, + PreserveOriginalFailure, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum Delivery { + InlineDirect, + InlineAwaited, + BlockingWorker, + BackgroundTask, + Deferred, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum Outcome { + Success, + Failure, + Abort, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +pub enum ErrorDisposition { + Preserve, + Replace, +} diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index 3868cd83b50..320f09ad478 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -1,201 +1,10 @@ +pub use crate::lifecycle::ocr::*; +pub use crate::lifecycle::{ErrorDisposition, Outcome}; + +#[cfg(test)] 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)] +use crate::ocr::{OcrRequest, prepare}; #[cfg(test)] mod tests { diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index d6e11c6e543..8b4ac9b92f5 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -38,7 +38,7 @@ fn request_config( Ok((provider, config)) } -pub(super) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> { +pub(crate) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> { check_admission_capabilities(request, &|key| std::env::var(key).ok()) } diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 6ab7bf39bb5..cb587d1c12e 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -4,7 +4,7 @@ litellm-python-bridge is the PyO3 cdylib exposing LiteLLM Rust APIs to the Pytho - Keep it thin: no business logic, no transforms, no I/O orchestration; just marshal in/out and call the core entrypoint - One stable method per top-level route (`ocr`/`aocr`, `messages`/`amessages`, ...); do not add per-provider helpers - Provider dispatch lives in `litellm-core`, never here -- Put domain-neutral Python/Serde conversion and GIL primitives in `litellm-python-interop` +- Put domain-neutral Python/Serde conversion, retained callbacks, and Python↔Tokio execution in `litellm-python-interop` - Two logical parts, kept as modules: the domain adapter (`src/routes/*`, `marshal`, `errors`) and the binding artifact (`#[pymodule]`, `#[pyfunction]`, registration in `lib.rs`) - Data handling: do not log OCR payloads or provider responses; avoid copying large payloads; sanitize errors before they cross the boundary - Tests: `cargo test --workspace` compiles here; Python tests cover disabled, enabled, and module-missing fallback for every exposed route diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bda09a7d840..8227419bdc9 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,7 +21,6 @@ trace-parity = [ ] [dependencies] -futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } @@ -30,10 +29,11 @@ pyo3.workspace = true pyo3-async-runtimes.workspace = true serde.workspace = true serde_json.workspace = true -tokio.workspace = true [dev-dependencies] criterion = "0.8.2" +futures-util.workspace = true +tokio.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index cc153a89b8f..00c3827eb9f 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,13 +1,4 @@ -use litellm_python_interop::release_count; use pyo3::prelude::*; -use pyo3::types::PyDict; - -#[pyfunction] -fn gil_stats(py: Python<'_>) -> PyResult> { - let stats = PyDict::new(py); - stats.set_item("releases", release_count())?; - Ok(stats.into_any().unbind()) -} #[cfg(feature = "panic-test")] #[pyfunction] @@ -15,9 +6,8 @@ fn _panic_for_test() { panic!("intentional PyO3 panic smoke test"); } -pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(gil_stats, module)?)?; +pub(crate) fn register(_module: &Bound<'_, PyModule>) -> PyResult<()> { #[cfg(feature = "panic-test")] - module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; + _module.add_function(wrap_pyfunction!(_panic_for_test, _module)?)?; Ok(()) } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 384f0be5a1b..64a5c5bda76 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,5 @@ mod diagnostics; mod errors; -mod execution; #[cfg(feature = "trace-parity")] mod function_trace; mod marshal; @@ -106,7 +105,6 @@ mod tests { "chat_completions", "achat_completions", "ResponsesWebSocketConnection", - "gil_stats", ]; let public_names: Vec = module diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 1e66a207592..45de81d8d80 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -31,7 +31,7 @@ macro_rules! bridge_route { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_sync(py, future, $map_error) + litellm_python_interop::run_sync(py, future, $map_error) } #[pyfunction] @@ -46,7 +46,7 @@ macro_rules! bridge_route { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_async(py, future, $map_error) + litellm_python_interop::run_async(py, future, $map_error) } pub(super) fn register( @@ -75,7 +75,7 @@ macro_rules! bridge_route { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_sync( + litellm_python_interop::run_sync( py, $crate::function_trace::capture(future), $map_error, @@ -94,7 +94,7 @@ macro_rules! bridge_route { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_async( + litellm_python_interop::run_async( py, $crate::function_trace::capture(future), $map_error, diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs index 97ff93f299a..8fc44bfeacd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs @@ -2,6 +2,7 @@ use pyo3::prelude::*; use serde_json::Value; use crate::errors::core_error_to_pyerr; +use litellm_python_interop::run_async; #[pyfunction] fn gateway_messages<'py>( @@ -17,7 +18,7 @@ fn gateway_messages<'py>( api_base, body, ); - crate::execution::run_async( + run_async( py, crate::function_trace::capture(future), core_error_to_pyerr, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index e626b6f3d2f..b83e48e438a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -8,18 +8,16 @@ use litellm_core::ocr::lifecycle::{ }; use litellm_core::ocr::types::{OcrDocumentProjection, OcrRequest, PreparedOcr}; use litellm_core::routing_utils::provider::get_custom_llm_provider; -use litellm_python_interop::{ - InvocationMode, InvocationOutcome, PreparedCall, Pythonized, from_py, to_py, -}; +use litellm_python_interop::{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, PyTuple}; +use pyo3::types::PyDict; use serde_json::Value; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async_value, run_sync_value}; +use litellm_python_interop::{run_async_value, run_sync_value}; #[pyclass] struct OcrState { @@ -305,32 +303,24 @@ fn invoke( 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), + let (method, awaiting) = match operation { + Operation::Setup => ("setup", false), + Operation::DeploymentPre => ("deployment_pre", true), + Operation::Prepare => ("prepare", false), + Operation::Send if asynchronous => ("send", true), + Operation::Send => ("send_sync", false), + Operation::DeploymentSuccess => ("deployment_success", true), + Operation::DeploymentFailure => ("deployment_failure", true), + Operation::SyncSuccess => ("sync_success", false), + Operation::AsyncSuccess => ("async_success", false), + Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false), + Operation::SyncFailure => ("sync_failure", false), + Operation::AsyncFailure => ("async_failure", true), + Operation::Restore => ("restore", false), 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)), - } + let value = host.getattr(py, method)?.call0(py)?; + Ok((awaiting, value)) } #[pyfunction] diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs deleted file mode 100644 index d397d20b9fd..00000000000 --- a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs +++ /dev/null @@ -1,49 +0,0 @@ -use std::fs; -use std::path::{Path, PathBuf}; - -const DISALLOWED_OUTSIDE_INTEROP: &[&str] = &[ - "py.import(\"json\")", - "pythonize::", - "serde_json::to_string", - "serde_json::from_str", -]; - -fn source_root() -> PathBuf { - Path::new(env!("CARGO_MANIFEST_DIR")).join("src") -} - -fn rust_sources(directory: &Path) -> Vec { - fs::read_dir(directory) - .expect("bridge source directory should be readable") - .map(|entry| { - entry - .expect("bridge source entry should be readable") - .path() - }) - .flat_map(|path| { - if path.is_dir() { - rust_sources(&path) - } else if path.extension().is_some_and(|extension| extension == "rs") { - vec![path] - } else { - Vec::new() - } - }) - .collect() -} - -#[test] -fn serialization_uses_the_interop_boundary() { - let root = source_root(); - - for path in rust_sources(&root) { - let source = fs::read_to_string(&path).expect("bridge source should be readable"); - for disallowed in DISALLOWED_OUTSIDE_INTEROP { - assert!( - !source.contains(disallowed), - "{} bypasses litellm-python-interop with `{disallowed}`", - path.display() - ); - } - } -} diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md index 07c7cffff46..d45a5cef688 100644 --- a/litellm-rust/crates/python-interop/AGENTS.md +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -1,5 +1,6 @@ litellm-python-interop is the domain-neutral PyO3 foundation. -- Owns generic Python/Serde conversion and interpreter primitives (`gil`, `marshal`) +- Owns generic Python/Serde conversion (`marshal`) and Python↔Tokio execution (`execution`) - Depends on PyO3 but no LiteLLM domain crate; no route types, no API registration, no cdylib +- `execution` is generic over the caller's error type (`map_error: fn(E) -> PyErr`); the bridge passes its own mapper - Keep it free of `litellm-core`, `OcrRequest`, `CallServices`, LiteLLM exceptions or `_native` surface diff --git a/litellm-rust/crates/python-interop/Cargo.toml b/litellm-rust/crates/python-interop/Cargo.toml index 5b8cf986a27..b0ff9bbb5ab 100644 --- a/litellm-rust/crates/python-interop/Cargo.toml +++ b/litellm-rust/crates/python-interop/Cargo.toml @@ -10,14 +10,13 @@ autotests = false name = "synthetic" path = "tests/synthetic/mod.rs" -[[test]] -name = "integration" -path = "tests/integration/mod.rs" - [dependencies] +futures-util.workspace = true pyo3.workspace = true +pyo3-async-runtimes.workspace = true pythonize.workspace = true serde.workspace = true +tokio.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/python-interop/src/callback.rs b/litellm-rust/crates/python-interop/src/callback.rs deleted file mode 100644 index 3942b150daf..00000000000 --- a/litellm-rust/crates/python-interop/src/callback.rs +++ /dev/null @@ -1,113 +0,0 @@ -use pyo3::prelude::*; -use pyo3::pyclass::{PyTraverseError, PyVisit}; -use pyo3::sync::PyOnceLock; -use pyo3::types::{PyDict, PyTuple}; - -use crate::constants::{ - AWAIT_ADAPTER_FILENAME, AWAIT_ADAPTER_FUNCTION, AWAIT_ADAPTER_MODULE, AWAIT_ADAPTER_SOURCE, -}; - -static AWAIT_ADAPTER: PyOnceLock> = PyOnceLock::new(); - -/// How a retained callback is bound to its caller. -/// -/// `Direct` mirrors `callable(*args, **kwargs)`: a coroutine returned by the -/// callable is handed back untouched and never awaited. `Await` mirrors -/// `await callable(*args, **kwargs)` inline in the caller's task. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum InvocationMode { - Direct, - Await, -} - -/// Result of [`PreparedCall::invoke`]. -/// -/// `Awaitable` carries an adapter coroutine that has not yet called the -/// callback. The callback runs, and any exception it raises surfaces, only -/// when Python drives that coroutine. -#[derive(Debug)] -pub enum InvocationOutcome { - Returned(Py), - Awaitable(Py), -} - -/// A callback plus its arguments, retained as owning Python references. -/// -/// Arguments are passed to the callback by identity, never copied, so the -/// callback observes and may mutate the caller's objects. Dropping the value -/// releases the references; a Python-visible owner must also expose them to -/// the cycle collector via [`PreparedCall::traverse`]. -pub struct PreparedCall { - mode: InvocationMode, - callable: Py, - positional: Py, - keywords: Option>, -} - -impl PreparedCall { - pub fn new( - mode: InvocationMode, - callable: Py, - positional: Py, - keywords: Option>, - ) -> Self { - Self { - mode, - callable, - positional, - keywords, - } - } - - pub fn invoke(&self, py: Python<'_>) -> PyResult { - match self.mode { - InvocationMode::Direct => self - .callable - .call( - py, - self.positional.bind(py), - self.keywords.as_ref().map(|kwargs| kwargs.bind(py)), - ) - .map(InvocationOutcome::Returned), - InvocationMode::Await => await_adapter(py)? - .call1(py, (&self.callable, &self.positional, &self.keywords)) - .map(InvocationOutcome::Awaitable), - } - } - - pub fn clone_ref(&self, py: Python<'_>) -> Self { - Self { - mode: self.mode, - callable: self.callable.clone_ref(py), - positional: self.positional.clone_ref(py), - keywords: self.keywords.as_ref().map(|value| value.clone_ref(py)), - } - } - - pub fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.callable)?; - visit.call(&self.positional)?; - if let Some(keywords) = &self.keywords { - visit.call(keywords)?; - } - Ok(()) - } -} - -/// Compiling the adapter runs Python, which may re-enter this function through -/// audit hooks. `PyOnceLock` forbids re-entrant initialization, so compile -/// first and only publish a finished adapter into the cell. -fn await_adapter(py: Python<'_>) -> PyResult<&Py> { - if let Some(adapter) = AWAIT_ADAPTER.get(py) { - return Ok(adapter); - } - let compiled = PyModule::from_code( - py, - AWAIT_ADAPTER_SOURCE, - AWAIT_ADAPTER_FILENAME, - AWAIT_ADAPTER_MODULE, - )? - .getattr(AWAIT_ADAPTER_FUNCTION)? - .unbind(); - Ok(AWAIT_ADAPTER.get_or_init(py, || compiled)) -} diff --git a/litellm-rust/crates/python-interop/src/constants.rs b/litellm-rust/crates/python-interop/src/constants.rs deleted file mode 100644 index d56e5babc86..00000000000 --- a/litellm-rust/crates/python-interop/src/constants.rs +++ /dev/null @@ -1,18 +0,0 @@ -use std::ffi::CStr; - -/// Python source of the coroutine adapter that awaits a retained callback -/// inline in the caller's task. It is compiled once per interpreter. -pub(crate) const AWAIT_ADAPTER_SOURCE: &CStr = - c"async def invoke_awaited(callable, positional, keywords): - if keywords is None: - return await callable(*positional) - return await callable(*positional, **keywords) -"; - -/// Filename recorded on the adapter's code object. Visible to Python -/// `compile` audit hooks and tracebacks. -pub const AWAIT_ADAPTER_FILENAME: &CStr = c"retained_callback.py"; - -pub(crate) const AWAIT_ADAPTER_MODULE: &CStr = c"_retained_callback"; - -pub(crate) const AWAIT_ADAPTER_FUNCTION: &str = "invoke_awaited"; diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-interop/src/execution.rs similarity index 84% rename from litellm-rust/crates/python-bridge/src/execution.rs rename to litellm-rust/crates/python-interop/src/execution.rs index ca0f8a4ae86..5928bc5355b 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-interop/src/execution.rs @@ -3,35 +3,34 @@ use std::panic::AssertUnwindSafe; use std::time::Duration; use futures_util::FutureExt; -use litellm_core::error::Error; -use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use serde::Serialize; use tokio::runtime::{Handle, Runtime}; use tokio::time::{self, MissedTickBehavior}; -pub(crate) fn run_sync( +use crate::marshal::Pythonized; +use crate::marshal::panic_to_pyerr; + +pub fn run_sync( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { let result = run_sync_value(py, future, map_error)?; Pythonized(result).into_pyobject(py).map(Bound::unbind) } -pub(crate) fn run_sync_value( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult +pub fn run_sync_value(py: Python<'_>, future: F, map_error: fn(E) -> PyErr) -> PyResult where T: Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { run_sync_value_on( py, @@ -41,15 +40,16 @@ where ) } -fn run_sync_value_on( +fn run_sync_value_on( py: Python<'_>, runtime: &Runtime, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult where T: Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { if Handle::try_current().is_ok() { return Err(PyRuntimeError::new_err( @@ -57,18 +57,19 @@ where )); } - let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; + let result = py.detach(move || runtime.block_on(wait_for_sync_result(future)))?; map_core_result(result, map_error) } -pub(crate) fn run_async( +pub fn run_async( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { pyo3_async_runtimes::tokio::future_into_py(py, async move { let result = catch_future_panic(future).await?; @@ -77,16 +78,20 @@ where }) } -pub(crate) async fn run_async_value(future: F, map_error: fn(Error) -> PyErr) -> PyResult +pub async fn run_async_value(future: F, map_error: fn(E) -> PyErr) -> PyResult where T: Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { let result = catch_future_panic(future).await?; map_core_result(result, map_error) } -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { +fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult +where + E: Send + 'static, +{ match result { Ok(value) => Ok(value), Err(error) => Err( @@ -96,9 +101,9 @@ fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) - } } -async fn catch_future_panic(future: F) -> PyResult> +async fn catch_future_panic(future: F) -> PyResult> where - F: Future>, + F: Future>, { AssertUnwindSafe(future) .catch_unwind() @@ -106,9 +111,9 @@ where .map_err(panic_to_pyerr) } -async fn wait_for_sync_result(future: F) -> PyResult> +async fn wait_for_sync_result(future: F) -> PyResult> where - F: Future>, + F: Future>, { let future = catch_future_panic(future); tokio::pin!(future); @@ -142,11 +147,18 @@ mod tests { use super::*; - fn runtime_error(error: Error) -> PyErr { - PyRuntimeError::new_err(error.to_string()) + #[derive(Debug)] + enum TestError { + InvalidRequest(&'static str), } - fn panicking_error_mapper(_error: Error) -> PyErr { + fn runtime_error(error: TestError) -> PyErr { + PyRuntimeError::new_err(match error { + TestError::InvalidRequest(message) => message, + }) + } + + fn panicking_error_mapper(_error: TestError) -> PyErr { panic!("error mapper panicked") } @@ -258,7 +270,7 @@ mod tests { let error = runtime.block_on(async { Python::attach(|py| { - run_sync::(py, async { Ok(true) }, runtime_error) + run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) }); @@ -294,9 +306,9 @@ mod tests { fn sync_runner_maps_a_panicked_future() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, - poll_fn(|_| -> Poll> { panic!("route future panicked") }), + poll_fn(|_| -> Poll> { panic!("route future panicked") }), runtime_error, ) .expect_err("panicked route should become a Python exception"); @@ -310,9 +322,9 @@ mod tests { fn sync_runner_maps_a_panicked_error_mapper() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, - async { Err(Error::InvalidRequest("invalid".to_string())) }, + async { Err(TestError::InvalidRequest("invalid")) }, panicking_error_mapper, ) .expect_err("panicked mapper should become a Python exception"); @@ -337,7 +349,7 @@ mod tests { #[test] fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { Python::initialize(); - let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let barrier = Arc::new(std::sync::Barrier::new(2)); let callers: Vec<_> = (0..2) .map(|_| { let barrier = Arc::clone(&barrier); @@ -348,9 +360,11 @@ mod tests { run_sync( py, async move { - Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait()) + let barrier = barrier; + tokio::task::spawn_blocking(move || barrier.wait()) .await - .is_ok()) + .map(|_| true) + .map_err(|_| TestError::InvalidRequest("join")) }, runtime_error, ), diff --git a/litellm-rust/crates/python-interop/src/gil.rs b/litellm-rust/crates/python-interop/src/gil.rs deleted file mode 100644 index 04b966a6002..00000000000 --- a/litellm-rust/crates/python-interop/src/gil.rs +++ /dev/null @@ -1,21 +0,0 @@ -use std::sync::atomic::{AtomicU64, Ordering}; - -use pyo3::prelude::*; - -static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); - -/// Runs work detached from the interpreter and records the release. -/// -/// `f` must not access Python state while the interpreter is detached. -pub fn release_gil(py: Python<'_>, f: F) -> T -where - F: FnOnce() -> T + Send, - T: Send, -{ - GIL_RELEASES.fetch_add(1, Ordering::Relaxed); - py.detach(f) -} - -pub fn release_count() -> u64 { - GIL_RELEASES.load(Ordering::Relaxed) -} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index 28c21d5ca3a..bd638ef8a73 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -1,9 +1,5 @@ -mod callback; -mod constants; -mod gil; +mod execution; mod marshal; -pub use callback::{InvocationMode, InvocationOutcome, PreparedCall}; -pub use constants::AWAIT_ADAPTER_FILENAME; -pub use gil::{release_count, release_gil}; +pub use execution::{run_async, run_async_value, run_sync, run_sync_value}; pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; diff --git a/litellm-rust/crates/python-interop/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs index a16d1e0ae13..46b7d04b54b 100644 --- a/litellm-rust/crates/python-interop/src/marshal.rs +++ b/litellm-rust/crates/python-interop/src/marshal.rs @@ -18,7 +18,8 @@ pub fn to_py(py: Python<'_>, value: &T) -> PyResult> where T: Serialize + ?Sized, { - pythonize::pythonize(py, value) + catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, value))) + .map_err(panic_to_pyerr)? .map(Bound::unbind) .map_err(|error| PyValueError::new_err(error.to_string())) } diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py deleted file mode 100644 index 59469a6545e..00000000000 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py +++ /dev/null @@ -1,848 +0,0 @@ -import asyncio -import atexit -import contextvars -import gc -import json -import threading -import weakref -from datetime import datetime -from unittest import TestCase - -import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import safe_deep_copy -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.logging_worker import LoggingWorker -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices, Usage - - -def logger_for( - callbacks=(), - stream=False, - input_callbacks=(), - sync_callbacks=(), - failure_callbacks=(), - async_failure_callbacks=(), - call_type="acompletion", -): - return Logging( - model="test", - messages=[{"role": "user", "content": "test"}], - stream=stream, - call_type=call_type, - start_time=datetime.now(), - litellm_call_id="retained-test", - function_id="retained-test", - dynamic_async_success_callbacks=list(callbacks), - dynamic_input_callbacks=list(input_callbacks), - dynamic_success_callbacks=list(sync_callbacks), - dynamic_failure_callbacks=list(failure_callbacks), - dynamic_async_failure_callbacks=list(async_failure_callbacks), - ) - - -def invoke_pre_call(owners, logger, additional): - owner = owners.prepare(logger.pre_call, (logger.messages, "test-key"), {"additional_args": additional}) - try: - return owner.invoke() - finally: - owner.close() - - -async def pre_call_identity_and_ignored_returns(owners): - saved = [] - ignored = {"replacement": True} - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - saved.append((kwargs, messages)) - return ignored - - logger = logger_for(input_callbacks=[Retain(), Retain()]) - details = logger.model_call_details - additional = {"headers": {"test": "header"}} - assert invoke_pre_call(owners, logger, additional) is None - assert len(saved) == 2 - assert saved[0][0] is saved[1][0] is details - assert saved[0][1] is saved[1][1] is logger.messages is details["input"] - assert details["additional_args"] is additional - assert "replacement" not in details - - -async def pre_call_mutations_visible_to_later_callbacks(owners): - saved, observed, order = [], [], [] - metadata = {"secret": "private", "keep": []} - removed = object() - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append("retain") - saved.append(kwargs) - - class Mutate(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append("mutate") - kwargs["normalized"] = "normalized" - assert kwargs.pop("remove") is removed - kwargs["retained_metadata"]["secret"] = "masked" - return {"replacement": True} - - class Observe(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append("observe") - observed.append((kwargs["normalized"], "remove" in kwargs, kwargs["retained_metadata"]["secret"])) - - logger = logger_for(input_callbacks=[Retain(), Mutate(), Observe()]) - details = logger.model_call_details - details.update(retained_metadata=metadata, normalized=None, remove=removed) - assert invoke_pre_call(owners, logger, {}) is None - assert order == ["retain", "mutate", "observe"] - assert observed == [("normalized", False, "masked")] - assert len(saved) == 1 and saved[0] is details - assert details["retained_metadata"] is metadata - assert metadata == {"secret": "masked", "keep": []} - assert details["normalized"] == "normalized" and "remove" not in details - assert "replacement" not in details - - -async def pre_call_mutation_survives_failure(owners): - observed, order = [], [] - metadata = {"keep": []} - lock = threading.Lock() - - class Fail(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append("fail") - kwargs["lock"] = lock - kwargs["retained_metadata"]["keep"].append("before failure") - raise RuntimeError("expected pre-call callback failure") - - class Observe(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append("observe") - observed.append((kwargs, tuple(kwargs["retained_metadata"]["keep"]), kwargs["lock"])) - - logger = logger_for(input_callbacks=[Fail(), Observe()]) - details = logger.model_call_details - details["retained_metadata"] = metadata - assert invoke_pre_call(owners, logger, {}) is None - assert order == ["fail", "observe"] - assert len(observed) == 1 and observed[0][0] is details - assert observed[0][1] == ("before failure",) and observed[0][2] is lock - assert details["retained_metadata"] is metadata - assert metadata == {"keep": ["before failure"]} and details["lock"] is lock - with TestCase().assertRaises(TypeError): - json.dumps({"lock": details["lock"]}) - - -async def real_post_call_logging(owners): - saved, observed, order = [], [], [] - shared = {"values": []} - response = ModelResponse(model="test") - ignored = {"replacement": True} - error = RuntimeError("expected post-call callback failure") - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - kwargs["stash"] = shared - saved.append(kwargs) - - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - order.append("retain") - saved.append(kwargs) - return ignored - - class MutateThenFail(CustomLogger): - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - order.append("fail") - kwargs["stash"]["values"].append("post") - kwargs["callback_error"] = error - kwargs["original_response"].choices[0].message.content = "mutated" - raise error - - class Observe(CustomLogger): - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - order.append("observe") - observed.append((kwargs, response_obj, start_time, end_time, tuple(kwargs["stash"]["values"]))) - return ignored - - logger = logger_for(input_callbacks=[Retain(), MutateThenFail(), Observe()]) - details = logger.model_call_details - assert invoke_pre_call(owners, logger, {}) is None - additional = {"headers": {"test": "post"}} - owner = owners.prepare(logger.post_call, (response, logger.messages, "test-key"), {"additional_args": additional}) - try: - assert owner.invoke() is None - finally: - owner.close() - assert order == ["retain", "fail", "observe"] - assert len(saved) == 2 and saved[0] is saved[1] is details - assert len(observed) == 1 and observed[0][0] is details - assert observed[0][1] is None and observed[0][2] is logger.start_time and observed[0][3] is None - assert observed[0][4] == ("post",) - assert details["original_response"] is response and response.choices[0].message.content == "mutated" - assert details["input"] is logger.messages and details["additional_args"] is additional - assert details["log_event_type"] == "post_api_call" and details["api_key"] == "test-key" - assert details["stash"] is shared and details["callback_error"] is error - assert "replacement" not in details - shared["values"].append("later") - assert saved[0]["stash"]["values"] == ["post", "later"] - - -async def real_post_call_dict_response(owners): - observed = [] - - class Observe(CustomLogger): - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - observed.append(kwargs["original_response"]) - - response = {"content": ["original"], "timestamp": datetime(2026, 1, 1)} - logger = logger_for(input_callbacks=[Observe()]) - owner = owners.prepare(logger.post_call, (response,)) - try: - assert owner.invoke() is None - finally: - owner.close() - assert len(observed) == 1 and observed[0] is logger.model_call_details["original_response"] - assert isinstance(observed[0], str) - assert json.loads(observed[0]) == {"content": ["original"], "timestamp": "2026-01-01 00:00:00"} - response["content"].append("later") - assert json.loads(observed[0])["content"] == ["original"] - - -async def real_sync_logging(owners): - saved, observations, replacements = [], [], [] - shared = {"values": []} - result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}]) - replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}]) - ignored = {"ignored": True}, result - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - kwargs["stash"] = shared - saved.append(kwargs) - - def logging_hook(self, kwargs, result, call_type): - observations.append(("retain", kwargs, result, call_type)) - kwargs["stash"]["values"].append("hook") - result.choices[0].message.content = "mutated" - return kwargs, result - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - kwargs["stash"]["values"].append("event") - observations.append(("event", kwargs, response_obj, tuple(kwargs["stash"]["values"]))) - return ignored - - class Replace(CustomLogger): - def logging_hook(self, kwargs, result, call_type): - observations.append(("replace", kwargs, result, result.choices[0].message.content)) - updated = {**kwargs, "adopted": True} - replacements.append(updated) - return updated, replacement - - class Observe(CustomLogger): - def logging_hook(self, kwargs, result, call_type): - observations.append(("observe", kwargs, result, tuple(kwargs["stash"]["values"]))) - return kwargs, result - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - observations.append(("success", kwargs, response_obj, tuple(kwargs["stash"]["values"]))) - - retain = Retain() - logger = logger_for(input_callbacks=[retain], sync_callbacks=[retain, Replace(), Observe()], call_type="completion") - assert invoke_pre_call(owners, logger, {}) is None - owner = owners.prepare(logger.success_handler, (result,)) - try: - assert owner.invoke() is None - finally: - owner.close() - assert [entry[0] for entry in observations] == ["retain", "replace", "observe", "event", "success"] - assert len(saved) == len(replacements) == 1 - assert observations[0][1] is observations[1][1] is saved[0] - assert observations[0][2] is observations[1][2] is result - assert observations[0][3] == "completion" and observations[1][3] == "mutated" - assert observations[2][3] == ("hook",) and observations[3][3] == observations[4][3] == ("hook", "event") - assert all(entry[1] is replacements[0] is logger.model_call_details for entry in observations[2:]) - assert all(entry[2] is replacement for entry in observations[2:]) - assert logger.model_call_details is not saved[0] - assert logger.model_call_details["adopted"] and "adopted" not in saved[0] - assert "ignored" not in logger.model_call_details - assert result.choices[0].message.content == "mutated" - assert replacement.choices[0].message.content == "replacement" - assert saved[0]["stash"] is logger.model_call_details["stash"] is shared - shared["values"].append("later") - assert observations[-1][1]["stash"]["values"] == ["hook", "event", "later"] - - -async def real_sync_logging_hook_failure(owners): - order, saved = [], [] - result = ModelResponse(model="test") - error = RuntimeError("expected sync logging hook failure") - - class MutateThenFail(CustomLogger): - def logging_hook(self, kwargs, result, call_type): - order.append("fail") - saved.append(kwargs) - kwargs["callback_error"] = error - result.choices[0].message.content = "before failure" - raise error - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - order.append("unexpected success") - - class Observe(CustomLogger): - def logging_hook(self, kwargs, result, call_type): - order.append("unexpected hook") - return kwargs, result - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - order.append("unexpected later success") - - logger = logger_for(sync_callbacks=[MutateThenFail(), Observe()], call_type="completion") - details = logger.model_call_details - owner = owners.prepare(logger.success_handler, (result,)) - try: - assert owner.invoke() is None - finally: - owner.close() - assert order == ["fail"] - assert len(saved) == 1 and saved[0] is logger.model_call_details is details - assert details["callback_error"] is error and error.__traceback__ is not None - assert result.choices[0].message.content == "before failure" - - -async def real_sync_failure_chain(owners): - await real_failure_chain(owners, awaited=False) - - -async def real_async_failure_chain(owners): - await real_failure_chain(owners, awaited=True) - - -async def real_failure_chain(owners, awaited): - saved, observations, hooks = [], [], [] - shared = {"values": []} - error = ValueError("provider failure") - callback_error = RuntimeError("expected failure callback error") - ignored = {"replacement": True}, object() - task = asyncio.current_task() - end = datetime.now() - - class Stage(CustomLogger): - def __init__(self, name): - super().__init__() - self.name = name - - def log_pre_api_call(self, model, messages, kwargs): - kwargs["stash"] = shared - saved.append(kwargs) - - def logging_hook(self, kwargs, result, call_type): - hooks.append("sync") - return ignored - - async def async_logging_hook(self, kwargs, result, call_type): - hooks.append("async") - return ignored - - def record(self, kwargs, response_obj, start_time, end_time): - observations.append( - ( - self.name, - kwargs, - response_obj, - kwargs["exception"], - tuple(kwargs["stash"]["values"]), - start_time, - end_time, - asyncio.current_task(), - ) - ) - if self.name == "fail": - kwargs["stash"]["values"].append("before failure") - kwargs["callback_error"] = callback_error - raise callback_error - return ignored - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - return self.record(kwargs, response_obj, start_time, end_time) - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - await asyncio.sleep(0) - return self.record(kwargs, response_obj, start_time, end_time) - - retain, fail, observe = Stage("retain"), Stage("fail"), Stage("observe") - callbacks = [retain, fail, observe] - logger = logger_for( - input_callbacks=[retain], - failure_callbacks=() if awaited else callbacks, - async_failure_callbacks=callbacks if awaited else (), - call_type="acompletion" if awaited else "completion", - ) - logger.model_call_details["litellm_params"]["acompletion"] = awaited - details = logger.model_call_details - assert invoke_pre_call(owners, logger, {}) is None - owner = owners.prepare( - logger.async_failure_handler if awaited else logger.failure_handler, - (error, "provider traceback"), - {"start_time": logger.start_time, "end_time": end}, - awaited=awaited, - ) - try: - if awaited: - assert await owner.invoke() is None - else: - assert owner.invoke() is None - finally: - owner.close() - assert [entry[0] for entry in observations] == ["retain", "fail", "observe"] - assert len(saved) == 1 and saved[0] is logger.model_call_details is details - assert all(entry[1] is details and entry[2] is None and entry[3] is error for entry in observations) - assert [entry[4] for entry in observations] == [(), (), ("before failure",)] - assert all(entry[5] is logger.start_time and entry[6] is end and entry[7] is task for entry in observations) - assert details["exception"] is error and details["callback_error"] is callback_error - assert callback_error.__traceback__ is not None - assert details["traceback_exception"] == "provider traceback" and details["log_event_type"] == "failed_api_call" - assert details["stash"] is shared and "replacement" not in details and hooks == [] - shared["values"].append("later") - assert saved[0]["stash"]["values"] == observations[-1][1]["stash"]["values"] == ["before failure", "later"] - - -async def real_async_logging(owners): - observations = [] - task = asyncio.current_task() - gate = asyncio.Event() - result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}]) - replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}]) - shared = {} - side_channel = {} - replaced_kwargs = [] - - class Retain(CustomLogger): - async def async_logging_hook(self, kwargs, result, call_type): - observations.append(("retained", kwargs, result, asyncio.current_task())) - return kwargs, result - - class MutateThenFail(CustomLogger): - async def async_logging_hook(self, kwargs, result, call_type): - asyncio.get_running_loop().call_soon(gate.set) - await gate.wait() - kwargs["retained_shared"]["changed"] = True - side_channel["failed_hook"] = kwargs - result.choices[0].message.content = "mutated" - raise RuntimeError("expected async callback failure") - - class Replace(CustomLogger): - async def async_logging_hook(self, kwargs, result, call_type): - observations.append(("replace", kwargs, result)) - updated = {**kwargs, "adopted": True} - replaced_kwargs.append(updated) - side_channel["replacement"] = replacement - return updated, replacement - - class Observe(CustomLogger): - async def async_logging_hook(self, kwargs, result, call_type): - observations.append(("observe", kwargs, result, asyncio.current_task())) - return kwargs, result - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - observations.append(("success", kwargs, response_obj)) - - logger = logger_for([Retain(), MutateThenFail(), Replace(), Observe()]) - logger.model_call_details["retained_shared"] = shared - owner = owners.prepare(logger.async_success_handler, (), {"result": result}, awaited=True) - await owner.invoke() - owner.close() - assert [entry[0] for entry in observations] == ["retained", "replace", "observe", "success"] - assert observations[0][3] is task and observations[2][3] is task - assert observations[0][2] is result and observations[1][2] is result - assert observations[2][2] is replacement and observations[3][2] is replacement - assert observations[0][1]["retained_shared"] is shared and shared["changed"] - assert observations[0][1] is observations[1][1] is side_channel["failed_hook"] - assert observations[2][1] is observations[3][1] is logger.model_call_details is replaced_kwargs[0] - assert logger.model_call_details is not observations[0][1] - assert logger.model_call_details["adopted"] and "adopted" not in observations[0][1] - assert logger.model_call_details["retained_shared"] is shared - assert side_channel["replacement"] is replacement - assert result.choices[0].message.content == "mutated" - observations[0][1]["retained_shared"]["after_replacement"] = True - assert observations[3][1]["retained_shared"]["after_replacement"] - - -async def real_copy_boundaries(owners): - lock = threading.Lock() - shared = {"values": []} - standard = { - "messages": [{"role": "user", "content": "private"}], - "response": {"choices": [{"message": {"content": "private"}}]}, - "metadata": shared, - } - details = {"standard_logging_object": standard, "shared": shared, "lock": lock} - passthrough = owners.prepare(CustomLogger().redact_standard_logging_payload_from_model_call_details, (details,)) - try: - assert passthrough.invoke() is details - finally: - passthrough.close() - logger = CustomLogger(turn_off_message_logging=True) - redact = owners.prepare(logger.redact_standard_logging_payload_from_model_call_details, (details,)) - try: - redacted = redact.invoke() - finally: - redact.close() - assert redacted is not details - assert redacted["standard_logging_object"] is not standard - assert redacted["shared"] is shared and redacted["lock"] is lock - assert redacted["standard_logging_object"]["metadata"] is shared - redacted["standard_logging_object"]["metadata"]["values"].append("shared mutation") - assert shared["values"] == ["shared mutation"] - assert redacted["standard_logging_object"]["messages"][0]["content"] == "redacted-by-litellm" - assert redacted["standard_logging_object"]["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" - assert standard["messages"][0]["content"] == "private" - assert standard["response"]["choices"][0]["message"]["content"] == "private" - - original_mode = litellm.safe_memory_mode - try: - for safe_mode in (False, True): - litellm.safe_memory_mode = safe_mode - uncopyable = {"lock": lock, "values": []} - values = [] - data = {"copyable": {"values": values, "alias": values}, "uncopyable": uncopyable} - owner = owners.prepare(safe_deep_copy, (data,)) - try: - copied = owner.invoke() - finally: - owner.close() - assert (copied is data) is safe_mode - assert (copied["copyable"] is data["copyable"]) is safe_mode - assert copied["copyable"]["values"] is copied["copyable"]["alias"] - assert (copied["copyable"]["values"] is values) is safe_mode - assert copied["uncopyable"] is uncopyable and copied["uncopyable"]["lock"] is lock - copied["copyable"]["values"].append("copy") - assert copied["copyable"]["alias"] == ["copy"] - copied["uncopyable"]["values"].append("fallback") - assert data["copyable"]["values"] == (["copy"] if safe_mode else []) - assert uncopyable["values"] == ["fallback"] - finally: - litellm.safe_memory_mode = original_mode - - -async def real_logging_worker(owners): - context = contextvars.ContextVar("component_worker_context", default="outside") - entered, release = asyncio.Event(), asyncio.Event() - observations = [] - worker = LoggingWorker(timeout=5, concurrency=1) - - class Payload: - pass - - async def upload(value, *, alias): - assert value is alias - assert context.get() == "submitted" - entered.set() - await release.wait() - observations.append((value.changed, context.get())) - context.set("worker only") - - payload = Payload() - payload.changed = False - reference = weakref.ref(payload) - invocation = owners.prepare(upload, (payload,), {"alias": payload}, awaited=True) - pending = invocation.invoke() - invocation.close() - enqueue = owners.prepare(worker.ensure_initialized_and_enqueue, (pending,)) - stop = owners.prepare(worker.stop, (), awaited=True) - flush = owners.prepare(worker.flush, (), awaited=True) - token = context.set("submitted") - try: - enqueue.invoke() - enqueue.close() - del pending, payload - context.set("consumer") - await entered.wait() - assert reference() is not None - reference().changed = True - release.set() - await flush.invoke() - assert observations == [(True, "submitted")] - assert context.get() == "consumer" - finally: - release.set() - enqueue.close() - flush.close() - await stop.invoke() - stop.close() - context.reset(token) - atexit.unregister(worker._flush_on_exit) - assert worker._worker_task is None and not worker._running_tasks and not worker._dequeued_tasks - assert worker._queue.empty() - gc.collect() - assert reference() is None - - -class ControlledStream: - def __init__(self): - self.originals = [ - ModelResponseStream(model="test", choices=[StreamingChoices(delta=Delta(content="hello"), index=0)]), - ModelResponseStream( - model="test", choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")] - ), - ModelResponseStream( - model="test", choices=[], usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8) - ), - ] - self.chunks = iter(self.originals) - self.closed = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self.chunks) - except StopIteration: - raise StopAsyncIteration - - async def aclose(self): - self.closed += 1 - - -async def real_stream_completion(owners): - logger = logger_for(stream=True) - completions = [] - cached = [] - cache_done = asyncio.Event() - - class CacheRecorder: - async def _add_streaming_response_to_cache(self, response): - cached.append(response) - cache_done.set() - - logger._llm_caching_handler = CacheRecorder() - - async def complete(response, cache_hit): - completions.append(response) - - logger._on_deferred_stream_complete = complete - stream = ControlledStream() - stream.originals[1].usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) - wrapper = CustomStreamWrapper( - completion_stream=stream, model="test", logging_obj=logger, custom_llm_provider="bedrock" - ) - pull = owners.prepare(wrapper.__anext__, (), awaited=True) - chunks = [] - retained_hidden = None - hidden_owner = None - while True: - try: - chunk = await pull.invoke() - chunks.append(chunk) - if len(chunks) == 1: - assert wrapper.chunks[-1] is chunk - chunk.choices[0].delta.content = "retained hello" - if chunk.choices and chunk.choices[0].finish_reason: - stored = wrapper.chunks[-1] - assert stored is not chunk and stored is not stream.originals[1] - assert stored.usage is stream.originals[1].usage - assert getattr(chunk, "usage", None) is None and stored.usage.total_tokens == 2 - retained_hidden = chunk._hidden_params - hidden_owner = owners.prepare(lambda value: value, (retained_hidden,)) - assert not completions - except StopAsyncIteration: - break - pull.close() - assert retained_hidden is not None - assert retained_hidden["usage"].total_tokens == 8 - assert completions == [] - response, cache_hit = logger._deferred_stream_complete_args - assert response.usage.total_tokens == 8 - assert response.choices[0].message.content == "retained hello" - assert retained_hidden["usage"] is response.usage - usage_chunk = wrapper.chunks[-1] - assert usage_chunk is not stream.originals[-1] - assert usage_chunk.usage is stream.originals[-1].usage - stream.originals[-1].usage.total_tokens = 13 - assert usage_chunk.usage.total_tokens == 13 - assert response.usage.total_tokens == 8 - deferred = owners.prepare(logger._on_deferred_stream_complete, (response, cache_hit), awaited=True) - logger._on_deferred_stream_complete = None - logger._deferred_stream_complete_args = None - close = owners.prepare(wrapper.aclose, (), awaited=True) - await close.invoke() - await close.invoke() - close.close() - assert stream.closed == 1 - del wrapper, logger - await deferred.invoke() - deferred.close() - assert completions == [response] - assert retained_hidden["usage"].total_tokens == 8 - assert hidden_owner.invoke() is retained_hidden - hidden_owner.close() - await cache_done.wait() - assert len(cached) == 1 and cached[0] is not response - assert cached[0].choices[0] is not response.choices[0] - cached[0].choices[0].message.content = "cache only" - assert response.choices[0].message.content == "retained hello" - - -async def real_sync_stream_copies(owners): - original_disable = litellm.disable_streaming_logging - copy_attempts = [] - - class Uncopyable: - def __deepcopy__(self, memo): - copy_attempts.append(True) - raise RuntimeError("expected streaming deepcopy failure") - - class CacheRecorder: - def __init__(self, responses): - self.responses = responses - - def _sync_add_streaming_response_to_cache(self, response): - self.responses.append(response) - - class Observe(CustomLogger): - def __init__(self, responses, finished): - super().__init__() - self.responses = responses - self.finished = finished - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - self.responses.append(response_obj) - self.finished.set() - - try: - litellm.disable_streaming_logging = True - for fallback in (False, True): - cached, logged = [], [] - finished = threading.Event() - - logger = logger_for(stream=True, sync_callbacks=[Observe(logged, finished)]) - logger._llm_caching_handler = CacheRecorder(cached) - source = ControlledStream() - source.originals[1].usage = source.originals[2].usage - wrapper = CustomStreamWrapper( - completion_stream=iter(source.originals[:2]), - model="test", - logging_obj=logger, - custom_llm_provider="bedrock", - ) - pull = owners.prepare(wrapper.__next__, ()) - close = owners.prepare(wrapper.aclose, (), awaited=True) - try: - first = pull.invoke() - assert wrapper.chunks[0] is first - first.choices[0].delta.content = "consumer mutation" - shared = {"values": []} - last = pull.invoke() - assert last.choices[0].finish_reason == "stop" - for chunk in wrapper.chunks: - chunk._hidden_params["retained_shared"] = shared - if fallback: - chunk._hidden_params["uncopyable"] = Uncopyable() - retained_hidden = last._hidden_params - with TestCase().assertRaises(StopIteration): - pull.invoke() - assert await asyncio.to_thread(finished.wait, 5) - finally: - pull.close() - await close.invoke() - close.close() - assert len(cached) == len(logged) == 1 - cache_response, log_response = cached[0], logged[0] - assert cache_response is not log_response - assert cache_response.choices[0].message.content == "consumer mutation" - assert log_response.choices[0].message.content == "consumer mutation" - assert retained_hidden["usage"].total_tokens == 8 - assert (cache_response.choices is log_response.choices) is fallback - assert (cache_response.usage is log_response.usage) is fallback - assert (cache_response.usage is retained_hidden["usage"]) is fallback - assert (cache_response._hidden_params is log_response._hidden_params) is fallback - assert (cache_response._hidden_params["retained_shared"] is shared) is fallback - cache_response.choices[0].message.content = "cache mutation" - cache_response._hidden_params["retained_shared"]["values"].append("cache mutation") - assert log_response.choices[0].message.content == ("cache mutation" if fallback else "consumer mutation") - assert shared["values"] == (["cache mutation"] if fallback else []) - assert log_response._hidden_params["retained_shared"]["values"] == (["cache mutation"] if fallback else []) - assert copy_attempts == [True] - finally: - litellm.disable_streaming_logging = original_disable - - -async def real_stream_close(owners): - source = ControlledStream() - logger = logger_for(stream=True) - wrapper = CustomStreamWrapper( - completion_stream=source, model="test", logging_obj=logger, custom_llm_provider="bedrock" - ) - pull = owners.prepare(wrapper.__anext__, (), awaited=True) - chunk = await pull.invoke() - assert wrapper.chunks[0] is chunk - pull.close() - retained = owners.prepare(lambda value: value, (chunk,)) - close = owners.prepare(wrapper.aclose, (), awaited=True) - await close.invoke() - await close.invoke() - close.close() - assert source.closed == 1 - assert wrapper.completion_stream is None - assert not getattr(logger, "_deferred_stream_complete_args", None) - assert retained.invoke() is chunk - assert chunk.choices[0].delta.content == "hello" - retained.close() - - -async def real_stream_cancellation(owners): - entered = asyncio.Event() - - class SuspendedStream(ControlledStream): - async def __anext__(self): - if self.chunks is not None: - chunk = next(self.chunks) - self.chunks = None - return chunk - entered.set() - await asyncio.Event().wait() - - source = SuspendedStream() - source.originals[0].usage = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) - logger = logger_for(stream=True) - wrapper = CustomStreamWrapper( - completion_stream=source, model="test", logging_obj=logger, custom_llm_provider="bedrock" - ) - pull = owners.prepare(wrapper.__anext__, (), awaited=True) - chunk = await pull.invoke() - retained = owners.prepare(lambda value: value, (chunk,)) - assert wrapper.chunks[0] is not chunk - assert wrapper.chunks[0].usage is source.originals[0].usage - task = asyncio.create_task(pull.invoke()) - pull.close() - await entered.wait() - task.cancel() - with TestCase().assertRaises(asyncio.CancelledError): - await task - assert logger.model_call_details.get("combined_usage_object") is None - recover = owners.prepare(wrapper._record_partial_usage_for_failure, ()) - try: - assert recover.invoke() is None - finally: - recover.close() - usage = logger.model_call_details["combined_usage_object"] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (3, 2, 5) - assert usage is not source.originals[0].usage - source.originals[0].usage.total_tokens = 99 - assert usage.total_tokens == 5 - retained_usage = owners.prepare(lambda value: value, (usage,)) - close = owners.prepare(wrapper.aclose, (), awaited=True) - await close.invoke() - await close.invoke() - close.close() - assert source.closed == 1 - assert wrapper.completion_stream is None and len(wrapper.chunks) == 1 - assert not getattr(logger, "_deferred_stream_complete_args", None) - assert retained.invoke() is chunk and chunk.choices[0].delta.content == "hello" - retained.close() - del wrapper, logger, usage - assert retained_usage.invoke().total_tokens == 5 - retained_usage.close() diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py deleted file mode 100644 index c841f444a13..00000000000 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py +++ /dev/null @@ -1,346 +0,0 @@ -import asyncio -import gc -import inspect -import weakref -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import Protocol, cast - -from callback_lifecycle import ReferenceFactory, run_checked, settle - - -class PreparedInvocation(Protocol): - def invoke(self) -> object: ... - - def close(self) -> None: ... - - -class CallFactory(Protocol): - def prepare( - self, - callable: Callable[..., object], - positional: tuple[object, ...], - keywords: dict[str, object] | None = None, - awaited: bool = False, - ) -> PreparedInvocation: ... - - -class LiveCallFactory(CallFactory, Protocol): - @property - def live(self) -> int: ... - - -def unchanged_result(value: object) -> object: - return value - - -@dataclass(frozen=True, slots=True) -class ResultTransformInvocation: - inner: PreparedInvocation - transform: Callable[[object], object] - awaited: bool - - def invoke(self) -> object: - result = self.inner.invoke() - transform = self.transform - if not self.awaited: - return transform(result) - - async def run() -> object: - return transform(await cast(Awaitable[object], result)) - - return run() - - def close(self) -> None: - self.inner.close() - - -@dataclass(frozen=True, slots=True) -class ResultTransformFactory: - inner: CallFactory - transform: Callable[[object], object] - - def prepare( - self, - callable: Callable[..., object], - positional: tuple[object, ...], - keywords: dict[str, object] | None = None, - awaited: bool = False, - ) -> PreparedInvocation: - return ResultTransformInvocation( - self.inner.prepare(callable, positional, keywords, awaited=awaited), self.transform, awaited - ) - - -@dataclass(frozen=True, slots=True) -class ExpiredBorrow: - edges: tuple[str, ...] - - -@dataclass(frozen=True, slots=True) -class CheckedWeakInvocation: - callback: weakref.ReferenceType[Callable[..., object]] - positional: tuple[weakref.ReferenceType[object], ...] - keywords: tuple[tuple[str, weakref.ReferenceType[object]], ...] - awaited: bool - - def resolve(self) -> tuple[Callable[..., object], tuple[object, ...], dict[str, object]] | ExpiredBorrow: - callback = self.callback() - positional = tuple(reference() for reference in self.positional) - keywords = {name: reference() for name, reference in self.keywords} - expired = ( - *(("callable",) if callback is None else ()), - *(f"positional:{index}" for index, value in enumerate(positional) if value is None), - *(f"keyword:{name}" for name, value in keywords.items() if value is None), - ) - if expired: - return ExpiredBorrow(expired) - assert callback is not None - return callback, positional, keywords - - def invoke(self) -> object: - if not self.awaited: - resolved = self.resolve() - if isinstance(resolved, ExpiredBorrow): - return resolved - callback, positional, keywords = resolved - return callback(*positional, **keywords) - - async def run() -> object: - resolved = self.resolve() - if isinstance(resolved, ExpiredBorrow): - return resolved - callback, positional, keywords = resolved - return await cast(Awaitable[object], callback(*positional, **keywords)) - - return run() - - def close(self) -> None: - pass - - -class CheckedWeakFactory: - def prepare( - self, - callable: Callable[..., object], - positional: tuple[object, ...], - keywords: dict[str, object] | None = None, - awaited: bool = False, - ) -> PreparedInvocation: - return CheckedWeakInvocation( - weakref.ref(callable), - tuple(weakref.ref(value) for value in positional), - tuple((name, weakref.ref(value)) for name, value in (keywords or {}).items()), - awaited, - ) - - -@dataclass(frozen=True, slots=True) -class MissingHandoffInvocation: - inner: PreparedInvocation - borrowed: PreparedInvocation - - def invoke(self) -> object: - return self.borrowed.invoke() - - def close(self) -> None: - try: - self.inner.close() - finally: - self.borrowed.close() - - -@dataclass(frozen=True, slots=True) -class MissingHandoffFactory: - inner: CallFactory - - def prepare( - self, - callable: Callable[..., object], - positional: tuple[object, ...], - keywords: dict[str, object] | None = None, - awaited: bool = False, - ) -> PreparedInvocation: - borrowed = CheckedWeakFactory().prepare(callable, positional, keywords, awaited=awaited) - return MissingHandoffInvocation(self.inner.prepare(callable, positional, keywords, awaited=awaited), borrowed) - - -def control_factory(control: str, inner: CallFactory) -> CallFactory: - if control == "identity": - return inner - if control == "weak": - return CheckedWeakFactory() - if control == "missing_handoff": - return MissingHandoffFactory(inner) - return {"result_passthrough": ResultTransformFactory(inner, unchanged_result)}[control] - - -@dataclass -class ControlNode: - stage: int = 0 - - -@dataclass -class LifetimeCallback: - awaited: bool - - def __call__(self, value: ControlNode, *, alias: ControlNode) -> object: - if self.awaited: - return self.run(value, alias=alias) - return value.stage + alias.stage - - async def run(self, value: ControlNode, *, alias: ControlNode) -> int: - await asyncio.sleep(0) - return value.stage + alias.stage - - -def prepare_released( - owners: CallFactory, awaited: bool -) -> tuple[ - PreparedInvocation, - tuple[ - weakref.ReferenceType[LifetimeCallback], weakref.ReferenceType[ControlNode], weakref.ReferenceType[ControlNode] - ], -]: - callback = LifetimeCallback(awaited) - value = ControlNode(13) - alias = ControlNode(29) - return ( - owners.prepare(callback, (value,), {"alias": alias}, awaited=awaited), - (weakref.ref(callback), weakref.ref(value), weakref.ref(alias)), - ) - - -@dataclass(frozen=True, slots=True) -class LifetimeObservation: - alive: tuple[bool, bool, bool] - result: int | ExpiredBorrow - - -async def deferred_lifetime(owners: CallFactory, awaited: bool) -> LifetimeObservation: - owner, references = prepare_released(owners, awaited) - try: - gc.collect() - alive = tuple(reference() is not None for reference in references) - pending = owner.invoke() - result = await settle(pending, awaited) - observation = LifetimeObservation(alive, result) - finally: - owner.close() - gc.collect() - assert all(reference() is None for reference in references) - return observation - - -@dataclass(frozen=True, slots=True) -class BorrowedObservation: - positional_identity: bool - keyword_identity: bool - positional_stage: int - keyword_stage: int - - -async def borrowed_lifetime(owners: CallFactory, awaited: bool) -> BorrowedObservation: - value, alias = ControlNode(), ControlNode() - value_ref, alias_ref = weakref.ref(value), weakref.ref(alias) - - def observe(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation: - return BorrowedObservation(value is value_ref(), alias is alias_ref(), value.stage, alias.stage) - - async def observe_async(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation: - return observe(value, alias=alias) - - owner = owners.prepare(observe_async if awaited else observe, (value,), {"alias": alias}, awaited=awaited) - try: - value.stage, alias.stage = 13, 29 - pending = owner.invoke() - value.stage, alias.stage = 17, 31 - return await settle(pending, awaited) - finally: - owner.close() - - -async def pending_handoff(owners: CallFactory, awaited: bool) -> LifetimeObservation: - assert awaited - owner, references = prepare_released(owners, awaited) - try: - pending = owner.invoke() - try: - owner.close() - gc.collect() - alive = tuple(reference() is not None for reference in references) - observation = LifetimeObservation(alive, await pending) - finally: - pending.close() - finally: - owner.close() - gc.collect() - assert all(reference() is None for reference in references) - return observation - - -async def direct_coroutine(owners: CallFactory, awaited: bool) -> bool: - async def body() -> int: - return 73 - - original = body() - try: - owner = owners.prepare(lambda: original, (), awaited=awaited) - try: - pending = owner.invoke() - if awaited: - assert await pending == 73 - nested = body() - try: - - async def returns_coroutine() -> Awaitable[int]: - return nested - - nested_owner = owners.prepare(returns_coroutine, (), awaited=True) - try: - result = await nested_owner.invoke() - return result is nested and inspect.getcoroutinestate(nested) == inspect.CORO_CREATED - finally: - nested_owner.close() - finally: - nested.close() - try: - return pending is original and inspect.getcoroutinestate(original) == inspect.CORO_CREATED - finally: - if inspect.iscoroutine(pending): - pending.close() - finally: - owner.close() - finally: - original.close() - - -def expected_control(witness: str, control: str, awaited: bool) -> object: - if witness in ("deferred_lifetime", "pending_handoff"): - if control == "weak" or (witness == "pending_handoff" and control == "missing_handoff"): - return LifetimeObservation( - (False, False, False), ExpiredBorrow(("callable", "positional:0", "keyword:alias")) - ) - return LifetimeObservation((True, True, True), 42) - if witness == "borrowed_lifetime": - return BorrowedObservation(True, True, 17 if awaited else 13, 31 if awaited else 29) - return {"direct_coroutine": True}[witness] - - -def run_control(witness: str, control: str, retained: bool, awaited: bool, factory: LiveCallFactory) -> None: - inner = factory if retained else ReferenceFactory() - owners = control_factory(control, inner) - - async def run() -> None: - observed = await WITNESSES[witness](owners, awaited) - assert observed == expected_control(witness, control, awaited), (witness, control, awaited, observed) - - run_checked(inner, run()) - - -WITNESSES: dict[str, Callable[[CallFactory, bool], object]] = { - "deferred_lifetime": deferred_lifetime, - "borrowed_lifetime": borrowed_lifetime, - "pending_handoff": pending_handoff, - "direct_coroutine": direct_coroutine, -} diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py deleted file mode 100644 index a085155b572..00000000000 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py +++ /dev/null @@ -1,667 +0,0 @@ -import asyncio -import copy -import gzip -import json -import threading -from collections import OrderedDict -from dataclasses import dataclass -from datetime import datetime -from typing import Literal -from unittest import TestCase - -import httpx -from fastapi import HTTPException - -import litellm -from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.datadog.datadog import DataDogLogger -from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger -from litellm.integrations.literal_ai import LiteralAILogger -from litellm.integrations.rubrik import RubrikLogger -from litellm.litellm_core_utils.litellm_logging import Logging, create_dummy_standard_logging_payload -from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler -from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail -from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import ModelResponse - - -async def integration_invoke(owners, callback, *args, **kwargs): - owner = owners.prepare(callback, args, kwargs, awaited=True) - pending = owner.invoke() - owner.close() - 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)) - - -@dataclass(frozen=True, slots=True) -class QueueObservation: - gcs_model_parameters: str - datadog_snapshot: str - literal_prepared_settings: str - - -async def real_logging_queue_copy_control(owners): - baseline = await integration_logging_queue_case(owners) - copied = await integration_logging_queue_case(owners, literal_copy="payload") - envelope = await integration_logging_queue_case(owners, literal_copy="envelope") - assert envelope == baseline - assert json.loads(baseline.gcs_model_parameters) == {"stream": True, "temperature": 0.25} - assert copied.datadog_snapshot == baseline.datadog_snapshot - assert copied.literal_prepared_settings == baseline.literal_prepared_settings - assert json.loads(copied.literal_prepared_settings) == {"stream": True} - assert json.loads(copied.gcs_model_parameters) == { - **json.loads(baseline.gcs_model_parameters), - "tools": [{"type": "function", "function": {"name": "lookup"}}], - } - return copied - - -def queue_loggers(entered, release, uploads): - class VertexTransport: - async def _ensure_access_token_async(self, **kwargs): - entered.set() - await release.wait() - return "fixture-token", "fixture-project" - - def _get_token_and_url(self, **kwargs): - return kwargs["auth_header"], None - - class Transport: - async def post(self, url, **kwargs): - wire = {**kwargs, "json": json.loads(json.dumps(kwargs["json"]))} if "json" in kwargs else kwargs - uploads.append((url, wire)) - return integration_response(url, {}, 202 if "datadog" in url else 200) - - datadog = DataDogLogger.__new__(DataDogLogger) - CustomBatchLogger.__init__(datadog, batch_size=100, flush_lock=asyncio.Lock()) - datadog.intake_url, datadog.DD_API_KEY, datadog.is_mock_mode = "https://datadog.invalid/logs", "test", False - datadog.async_client = Transport() - gcs = GCSBucketLogger.__new__(GCSBucketLogger) - CustomBatchLogger.__init__(gcs, batch_size=100) - gcs.log_queue = asyncio.Queue() - gcs.BUCKET_NAME, gcs.path_service_account_json = "fixture-bucket", None - gcs.vertex_instances = {"IAM_AUTH": VertexTransport()} - gcs.use_batched_logging = True - gcs.async_httpx_client = Transport() - literal = LiteralAILogger.__new__(LiteralAILogger) - CustomBatchLogger.__init__(literal, batch_size=100, flush_lock=asyncio.Lock()) - literal.literalai_api_url, literal.headers = "https://literal.invalid", {} - literal.async_httpx_client = Transport() - return datadog, gcs, literal - - -async def integration_logging_queue_case(owners, *, literal_copy: Literal["direct", "envelope", "payload"] = "direct"): - entered, release = asyncio.Event(), asyncio.Event() - uploads = [] - datadog, gcs, literal = queue_loggers(entered, release, uploads) - copy_payload = literal_copy == "payload" - - class LiteralCallback(CustomLogger): - def __init__(self, delegate): - super().__init__() - self.delegate = delegate - self.calls = self.completed = 0 - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - self.calls += 1 - self.received = kwargs - self.forwarded = { - **kwargs, - "standard_logging_object": ( - copy.deepcopy(kwargs["standard_logging_object"]) - if copy_payload - else kwargs["standard_logging_object"] - ), - } - await self.delegate.async_log_failure_event(self.forwarded, response_obj, start_time, end_time) - self.completed += 1 - - literal_callback = literal if literal_copy == "direct" else LiteralCallback(literal) - - payload = create_dummy_standard_logging_payload() - payload.update(status="failure", error_str="x" * 10001) - messages, settings, metadata = payload["messages"], payload["model_parameters"], payload["metadata"] - completion = payload["response"]["choices"][0]["message"] - tools = [{"type": "function", "function": {"name": "lookup"}}] - settings["tools"] = tools - now = datetime.now() - logging = Logging( - model="fixture-model", - messages=messages, - stream=False, - call_type="acompletion", - start_time=now, - litellm_call_id="fixture-queue", - function_id="fixture", - dynamic_async_failure_callbacks=[datadog, gcs, literal_callback], - ) - error = RuntimeError("fixture failure") - kwargs = logging.model_call_details - kwargs.update(standard_logging_object=payload, model="fixture-model", exception=error, end_time=now) - await integration_invoke(owners, logging.async_failure_handler, error, "fixture traceback", now, now) - if literal_copy != "direct": - assert literal_callback.calls == literal_callback.completed == 1 - assert literal_callback.received is kwargs and literal_callback.forwarded is not kwargs - assert (literal_callback.forwarded["standard_logging_object"] is payload) is (not copy_payload) - assert len(datadog.log_queue) == gcs.log_queue.qsize() == len(literal.log_queue) == 1 - assert kwargs["standard_logging_object"] is payload - assert payload["messages"] is messages and payload["model_parameters"] is settings - assert payload["error_str"].endswith("truncated by litellm, this logger does not support large content") - assert ("tools" in settings) is copy_payload - dd_snapshot = json.loads(datadog.log_queue[0]["message"]) - assert dd_snapshot["model_parameters"]["tools"] == tools - queued = gcs.log_queue.get_nowait() - assert queued["payload"] is payload and queued["kwargs"] is kwargs and queued["response_obj"] is None - gcs.log_queue.put_nowait(queued) - generation = literal.log_queue[0]["generation"] - prepared_settings = json.dumps(generation["settings"], sort_keys=True) - assert "tools" not in generation["settings"] and generation["tools"] == tools - if copy_payload: - assert generation["settings"] is not settings and generation["tools"] is not tools - assert generation["messages"] is not messages and generation["messageCompletion"] is not completion - assert literal.log_queue[0]["metadata"] is not metadata - else: - assert generation["settings"] is settings and generation["tools"] is tools - assert generation["messages"] is messages and generation["messageCompletion"] is completion - assert literal.log_queue[0]["metadata"] is metadata - - flush = asyncio.create_task(integration_invoke(owners, gcs.flush_queue)) - try: - await entered.wait() - assert not uploads and not flush.done() and gcs.log_queue.empty() - messages[0]["content"] = "mutated before serialization" - settings["temperature"] = 0.25 - completion["content"] = "late completion" - payload["messages"] = [{"role": "user", "content": "replacement field"}] - kwargs["standard_logging_object"] = {"replacement": True} - release.set() - await flush - finally: - release.set() - if not flush.done(): - flush.cancel() - await asyncio.gather(flush, return_exceptions=True) - assert len(uploads) == 1 - gcs_snapshot = json.loads(uploads[0][1]["data"]) - assert gcs_snapshot["messages"] == payload["messages"] - assert gcs_snapshot["model_parameters"] == settings - assert "replacement" not in gcs_snapshot - await integration_invoke(owners, datadog.flush_queue) - await integration_invoke(owners, literal.flush_queue) - assert len(uploads) == 3 and not datadog.log_queue and not literal.log_queue - sent_dd = json.loads(gzip.decompress(uploads[1][1]["data"])) - assert json.loads(sent_dd[0]["message"]) == dd_snapshot - literal_wire = uploads[2][1]["json"] - sent_generation = literal_wire["variables"]["generation_0"] - assert sent_generation["messages"] != gcs_snapshot["messages"] - if copy_payload: - assert sent_generation["messages"] == dd_snapshot["messages"] - assert json.dumps(sent_generation["settings"], sort_keys=True) == prepared_settings - assert sent_generation["messageCompletion"] == dd_snapshot["response"]["choices"][0]["message"] - else: - assert sent_generation["messages"] == messages - assert sent_generation["settings"]["temperature"] == 0.25 - assert sent_generation["messageCompletion"]["content"] == "late completion" - messages[0]["content"] = "after serialization" - assert sent_generation["messages"][0]["content"] == ( - "Hello, world!" if copy_payload else "mutated before serialization" - ) - assert dd_snapshot["messages"][0]["content"] == "Hello, world!" - return QueueObservation( - gcs_model_parameters=json.dumps(gcs_snapshot["model_parameters"], sort_keys=True), - datadog_snapshot=json.dumps(dd_snapshot, sort_keys=True), - literal_prepared_settings=prepared_settings, - ) - - -async def real_crowdstrike_translator_identity(owners): - entered, release = asyncio.Event(), asyncio.Event() - calls = [] - - class Transport: - async def post(self, url, json, **kwargs): - calls.append(json) - entered.set() - await release.wait() - return integration_response( - url, - { - "result": { - "blocked": False, - "transformed": True, - "guard_output": {"messages": [{"role": "user", "content": "redacted"}]}, - } - }, - ) - - guardrail = CrowdStrikeAIDRHandler.__new__(CrowdStrikeAIDRHandler) - CustomGuardrail.__init__(guardrail, guardrail_name="fixture-crowdstrike", event_hook=GuardrailEventHooks.pre_call) - guardrail.api_base, guardrail.api_key, guardrail.fail_on_error = "https://crowdstrike.invalid", "test", True - guardrail.skip_system_message_in_guardrail = True - guardrail.async_handler = Transport() - system = {"role": "system", "content": "internal policy"} - user = {"role": "user", "content": "private text", "extra": {"retained": True}} - messages = [system, user] - data = {"model": "fixture-model", "messages": messages} - task = asyncio.create_task( - integration_invoke( - owners, - OpenAIChatCompletionsHandler().process_input_messages, - data, - guardrail, - ) - ) - try: - await entered.wait() - assert data["messages"] is messages and not task.done() - assert calls[0]["guard_input"]["messages"] == [{"role": "user", "content": "private text"}] - user["extra"]["during_http"] = True - release.set() - assert await task is data - finally: - release.set() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - assert data["messages"] is not messages - assert data["messages"][0] is system - assert data["messages"][1] is not user - assert data["messages"][1]["extra"] is user["extra"] - assert data["messages"][1]["content"] == "redacted" - assert user["content"] == "private text" and messages[1] is user - detached = copy.deepcopy(messages[1:]) - inputs = {"texts": ["private text"], "structured_messages": detached} - control = await integration_invoke(owners, guardrail.apply_guardrail, inputs, {"messages": messages}, "request") - assert len(calls) == 2 - assert control["structured_messages"] is detached and detached[0] is not user - assert control["texts"] == ["redacted"] - assert detached[0]["content"] == user["content"] == "private text" - - -async def real_rubrik_block_lifecycle(owners): - for input_type, populated in (("request", False), ("response", True)): - entered, release = asyncio.Event(), asyncio.Event() - moderation, uploads = [], [] - - class Transport: - def __init__(self, moderation, uploads, entered, release): - self.moderation, self.uploads = moderation, uploads - self.entered, self.release = entered, release - - async def post(self, url, json, **kwargs): - if url.endswith("/batch"): - self.uploads.append(json) - return integration_response(url, {}) - self.moderation.append(json) - self.entered.set() - await self.release.wait() - return integration_response(url, {"choices": [{"message": {"content": "blocked by policy"}}]}) - - rubrik = RubrikLogger.__new__(RubrikLogger) - CustomGuardrail.__init__( - rubrik, - guardrail_name="fixture-rubrik", - event_hook=GuardrailEventHooks.post_call, - flush_lock=asyncio.Lock(), - batch_size=100, - ) - rubrik._periodic_flush_task = None - rubrik.sampling_rate, rubrik._headers = 1.0, {} - rubrik._dropped_since_warning, rubrik._last_drop_warning_time = 0, 0.0 - rubrik.prompt_moderation_endpoint = "https://rubrik.invalid/before" - rubrik.response_moderation_endpoint = "https://rubrik.invalid/after" - rubrik.logging_endpoint = "https://rubrik.invalid/batch" - rubrik.moderation_client = rubrik.async_httpx_client = Transport(moderation, uploads, entered, release) - other = RubrikLogger.__new__(RubrikLogger) - CustomGuardrail.__init__(other, guardrail_name="fixture-other-rubrik", event_hook=GuardrailEventHooks.post_call) - logging = Logging( - model="fixture-model", - messages=[{"role": "user", "content": "original prompt"}], - stream=False, - call_type="acompletion", - start_time=datetime.now(), - litellm_call_id="fixture-correlation", - function_id="fixture", - ) - details = logging.model_call_details - details.update(messages=logging.messages, model="fixture-model", litellm_call_id="fixture-correlation") - details["system"] = "system scaffold" - if populated: - details["standard_logging_object"] = create_dummy_standard_logging_payload() - messages = details["messages"] - request = {"model": "fixture-model", "litellm_call_id": "fixture-correlation", "messages": messages} - inputs = {"texts": ["original response"], "structured_messages": messages} - success = owners.prepare(rubrik.async_log_success_event, (details, None, None, None), awaited=True) - task = asyncio.create_task( - integration_invoke(owners, rubrik.apply_guardrail, inputs, request, input_type, logging) - ) - try: - await entered.wait() - assert not task.done() and "_rubrik_logging_obj" not in request - assert "_rubrik_blocked" not in details - if input_type == "request": - assert moderation[0]["correlation_key"] == "fixture-correlation" - assert moderation[0]["messages"][0]["content"] == "original prompt" - else: - assert moderation[0]["request"]["messages"] is messages - assert moderation[0]["response"]["id"] == "fixture-correlation" - release.set() - with TestCase().assertRaises(ModifyResponseException) as caught: - await task - error = caught.exception - assert error.request_data is request and error.message == "blocked by policy" - assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True - await integration_invoke( - owners, other.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user") - ) - assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True - assert not other.log_queue and not rubrik.log_queue and not uploads - await integration_invoke( - owners, rubrik.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user") - ) - assert "_rubrik_logging_obj" not in request and details["_rubrik_blocked"] is True - assert len(rubrik.log_queue) == 1 - queued = rubrik.log_queue[0] - assert queued["id"] == "fixture-correlation" - assert queued["response"] == "ModifyResponseException: blocked by policy" - assert queued["messages"][0] == {"role": "system", "content": "system scaffold"} - assert messages[0] == {"role": "user", "content": "original prompt"} - if populated: - base = details["standard_logging_object"] - assert queued["metadata"] is not base["metadata"] - assert queued["messages"][1] is not base["messages"][0] - assert isinstance(base["response"], dict) - else: - assert queued["messages"][1] is messages[0] - assert queued["status"] == "failure" - assert queued["metadata"]["user_api_key_user_id"] == "fixture-user" - pending = success.invoke() - success.close() - assert await pending is None - assert len(rubrik.log_queue) == 1 and rubrik.log_queue[0] is queued - await integration_invoke(owners, rubrik.flush_queue) - assert not rubrik.log_queue and len(uploads) == 1 and uploads[0][0] is queued - finally: - success.close() - release.set() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - await rubrik.aclose() - if rubrik._periodic_flush_task is not None: - await asyncio.gather(rubrik._periodic_flush_task, return_exceptions=True) - - -async def real_parallel_guardrail_snapshots(owners): - original_mode = litellm.safe_memory_mode - try: - 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, *, 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 - - class Uncopyable: - def __init__(self): - self.attempts = 0 - self.observed = [] - - def __deepcopy__(self, memo): - self.attempts += 1 - raise TypeError("fixture cannot be copied") - - arrived, release, mutated = asyncio.Event(), asyncio.Event(), asyncio.Event() - observations = {} - sentinel = Uncopyable() - live = {"messages": [{"role": "user", "content": "original"}], "uncopyable": sentinel} - raw = independent_snapshot(live) - 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): - observations[self.guardrail_name] = data - if len(observations) == 4: - arrived.set() - await release.wait() - if self.guardrail_name == "fixture-live-writer": - data["messages"][0]["content"] = "shared mutation" - mutated.set() - await mutated.wait() - if self.scan_raw_request: - assert data["messages"][0]["content"] == "original" - data["messages"][0]["content"] = self.guardrail_name - assert data["uncopyable"] is sentinel - data["uncopyable"].observed.append(self.guardrail_name) - else: - 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( - Inspect( - guardrail_name=name, - event_hook=GuardrailEventHooks.pre_call, - default_on=True, - run_in_parallel=True, - scan_raw_request="raw" in name, - ) - for name in ("fixture-live-writer", "fixture-live-reader", "fixture-raw-a", "fixture-raw-b") - ) - 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, - dispatched, - raw, - UserAPIKeyAuth(), - "acompletion", - ) - ) - try: - await arrived.wait() - assert not task.done() - 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 + copy_live and not sentinel.observed - release.set() - 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 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 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(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() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - -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() - - class Transport: - async def post(self, url, **kwargs): - workers.append(threading.current_thread()) - calls.append((url, kwargs)) - if url.endswith("/token"): - entered.set() - 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"): - return integration_response(url, {}, headers={"etag": "fixture-etag"}) - assert url.endswith("/processContent") - assert kwargs["headers"]["If-None-Match"] == "fixture-etag" - return integration_response( - url, {"policyActions": [{"action": "restrictAccess", "restrictionAction": "block"}]} - ) - - purview = MicrosoftPurviewDLPGuardrail.__new__(MicrosoftPurviewDLPGuardrail) - CustomGuardrail.__init__(purview, guardrail_name="fixture-purview", event_hook=GuardrailEventHooks.logging_only) - purview.async_handler = Transport() - purview.tenant_id, purview.client_id, purview.client_secret = "fixture-tenant", "fixture-client", "test" - purview.purview_app_name, purview.user_id_field, purview.guardrail_provider = ( - "fixture", - "user_id", - "microsoft_purview", - ) - purview._token_cache, purview._scope_cache = None, OrderedDict() - purview._scope_cache_maxsize, purview._cache_lock = 1000, threading.Lock() - metadata = {"user_api_key_user_id": "fixture-user"} - kwargs = { - "messages": [{"role": "user", "content": "prompt at dispatch"}], - "litellm_params": {"metadata": metadata}, - "litellm_call_id": "before-http", - } - 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: - 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 - 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() - 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"] - assert entries[0]["contentEntries"][0]["content"]["data"] == "prompt at dispatch" - assert entries[1]["contentEntries"][0]["content"]["data"] == "response mutated while audit waits" - assert all(entry["contentEntries"][0]["correlationId"] == "after-http" for entry in entries) - assert kwargs["litellm_params"]["metadata"] is metadata - info = kwargs["metadata"]["standard_logging_guardrail_information"] - assert len(info) == 2 and all(item["guardrail_status"] == "guardrail_intervened" for item in info) - assert all(item["end_time"] >= item["start_time"] and item["duration"] >= 0 for item in info) - finally: - owner.close() - release.set() - 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 deleted file mode 100644 index 118ca5f3013..00000000000 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ /dev/null @@ -1,633 +0,0 @@ -import asyncio -import contextvars -import copy -import gc -import inspect -import json -import sys -import threading -import weakref -from unittest import TestCase - - -async def checkpoint(): - ready = asyncio.Event() - asyncio.get_running_loop().call_soon(ready.set) - await ready.wait() - - -async def settle(pending, awaited): - return await pending if awaited else pending - - -class Value: - pass - - -def cold_awaited_adapter_reentry(factory): - events = [] - compilations = [] - - def invoke_failure(): - error = LookupError("cold-cache callback error") - - async def failing(): - raise error - - owner = factory.prepare(failing, (), awaited=True) - pending = owner.invoke() - try: - assert inspect.getcoroutinestate(pending) == inspect.CORO_CREATED - with TestCase().assertRaises(LookupError) as caught: - pending.send(None) - assert caught.exception is error - finally: - pending.close() - owner.close() - error.__traceback__ = None - - def audit(event, args): - if event != "compile" or args[1] != "retained_callback.py": - return - compilations.append(args[1]) - if len(compilations) == 1: - events.append("entered") - invoke_failure() - events.append("nested completed") - - sys.addaudithook(audit) - invoke_failure() - events.append("outer completed") - assert events == ["entered", "nested completed", "outer completed"] - assert len(compilations) == 2 - assert factory.live == 0 - - result = object() - direct = factory.prepare(lambda value: value, (result,)) - try: - assert direct.invoke() is result - finally: - direct.close() - - async def successful(value, *, alias): - assert value is alias - await checkpoint() - return value - - owner = factory.prepare(successful, (result,), {"alias": result}, awaited=True) - try: - assert asyncio.run(owner.invoke()) is result - finally: - owner.close() - assert len(compilations) == 2 - assert factory.live == 0 - - -class ReferenceFactory: - def __init__(self): - self.live = 0 - - def prepare(self, callable, positional, keywords=None, awaited=False): - return ReferenceOwner(self, callable, positional, keywords, awaited) - - -class ReferenceOwner: - def __init__(self, factory, callable, positional, keywords, awaited): - self.factory = factory - self.call = (callable, positional, keywords, awaited) - factory.live += 1 - - def invoke(self): - if self.call is None: - raise RuntimeError("invocation owner released") - callable, positional, keywords, awaited = self.call - if not awaited: - return callable(*positional, **(keywords if keywords is not None else {})) - - async def run(): - return await callable(*positional, **(keywords if keywords is not None else {})) - - return run() - - def clone_owner(self): - return ReferenceOwner(self.factory, *self.call) - - def close(self): - if self.call is not None: - released, self.call = self.call, None - self.factory.live -= 1 - del released - - def __del__(self): - self.close() - - -async def awaitable_kinds(owners): - payload = Value() - calls = [] - - async def coroutine(value, *, alias): - assert value is alias - calls.append("called") - return value - - class CustomAwaitable: - def __await__(self): - return coroutine(payload, alias=payload).__await__() - - for kind in ("async", "sync_coroutine", "custom", "future"): - future = asyncio.get_running_loop().create_future() - future.set_result(payload) - callback = { - "async": coroutine, - "sync_coroutine": lambda value, *, alias: coroutine(value, alias=alias), - "custom": lambda value, *, alias: CustomAwaitable(), - "future": lambda value, *, alias, future=future: future, - }[kind] - owner = owners.prepare(callback, (payload,), {"alias": payload}, awaited=True) - pending = owner.invoke() - before = len(calls) - owner.close() - assert await pending is payload - assert len(calls) == before + (kind != "future") - - owner = owners.prepare(lambda: payload, (), awaited=True) - with TestCase().assertRaises(TypeError): - await owner.invoke() - owner.close() - - inner = coroutine(payload, alias=payload) - - async def returns_coroutine(): - return inner - - owner = owners.prepare(returns_coroutine, (), awaited=True) - assert await owner.invoke() is inner - assert inner.cr_frame is not None - inner.close() - owner.close() - - direct = owners.prepare(coroutine, (payload,), {"alias": payload}) - untouched = direct.invoke() - assert untouched.cr_frame is not None - assert untouched.cr_await is None - untouched.close() - direct.close() - - -async def identity_and_context(owners): - context = contextvars.ContextVar("retained_context", default="outside") - task = asyncio.current_task() - loop = asyncio.get_running_loop() - thread = threading.get_ident() - payload = {"nested": {}} - alias = payload["nested"] - saved = [] - gate = asyncio.Event() - - async def nested(value): - assert asyncio.current_task() is task - assert context.get() == "inside" - value["nested"]["nested_call"] = True - context.set("nested") - return value - - async def callback(value, *, shared): - assert value is payload and shared is alias - assert asyncio.current_task() is task - assert asyncio.get_running_loop() is loop - assert threading.get_ident() == thread - assert context.get() == "at_await" - owner.close() - payload["closure_mutation"] = True - context.set("inside") - saved.append(value) - shared["before"] = True - loop.call_soon(gate.set) - await gate.wait() - inner = owners.prepare(nested, (value,), awaited=True) - try: - assert await inner.invoke() is value - finally: - inner.close() - return value - - owner = owners.prepare(callback, (payload,), {"shared": alias}, awaited=True) - pending = owner.invoke() - context.set("at_await") - assert await pending is payload - assert context.get() == "nested" - assert payload["closure_mutation"] is True - assert owners.live == 0 - alias["after"] = True - assert saved[0]["nested"] == {"before": True, "nested_call": True, "after": True} - - -async def exceptions(owners): - for error in (RuntimeError("original"), KeyboardInterrupt("original"), asyncio.CancelledError("original")): - cause = ValueError("cause") - payload = {} - - async def callback(payload=payload, error=error, cause=cause): - payload["changed"] = True - raise error from cause - - owner = owners.prepare(callback, (), awaited=True) - caught_error = None - try: - await owner.invoke() - except BaseException as caught: - caught_error = caught - finally: - owner.close() - assert caught_error is error and caught_error.__cause__ is cause - frames = [] - tb = caught_error.__traceback__ - while tb: - frames.append(tb.tb_frame.f_code.co_name) - tb = tb.tb_next - assert "callback" in frames - assert payload["changed"] is True - - -async def exception_ownership(owners): - class Callback: - def __init__(self, error): - self.error = error - - async def __call__(self, value): - value.changed = True - raise self.error - - value = Value() - error = RuntimeError("retained exception") - callback = Callback(error) - value_ref, callback_ref = weakref.ref(value), weakref.ref(callback) - owner = owners.prepare(callback, (value,), awaited=True) - del value, callback - caught_error = None - try: - await owner.invoke() - except RuntimeError as caught: - caught_error = caught - assert caught_error is error - owner.close() - assert value_ref().changed and callback_ref() is not None - del error, caught_error - gc.collect() - assert value_ref() is None and callback_ref() is None - - -async def cancellation_before_start(owners): - started = [] - - async def callback(value): - started.append(value) - - for operation in ("close", "cancel"): - value = Value() - ref = weakref.ref(value) - owner = owners.prepare(callback, (value,), awaited=True) - pending = owner.invoke() - owner.close() - del value - assert ref() is not None - if operation == "close": - pending.close() - else: - task = asyncio.create_task(pending) - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - del task - del pending - await checkpoint() - gc.collect() - assert ref() is None - assert started == [] - - -async def cancellation_case(owners, repeated=False, suppress=False): - started, cleaning, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() - value = Value() - ref = weakref.ref(value) - observed = [] - - async def callback(argument): - try: - started.set() - await asyncio.Event().wait() - except asyncio.CancelledError: - cleaning.set() - try: - await finish.wait() - except asyncio.CancelledError: - observed.append("second cancellation") - await finish.wait() - argument.cleaned = True - if suppress: - return argument - raise - - owner = owners.prepare(callback, (value,), awaited=True) - task = asyncio.create_task(owner.invoke()) - owner.close() - del value - try: - await started.wait() - task.cancel() - await cleaning.wait() - assert ref() is not None and not task.done() - if repeated: - task.cancel() - barrier = asyncio.Event() - asyncio.get_running_loop().call_soon(barrier.set) - await barrier.wait() - assert observed == ["second cancellation"] and not task.done() - finish.set() - try: - result = await task - assert suppress and result is ref() and result.cleaned - del result - except asyncio.CancelledError: - assert not suppress - assert ref().cleaned - finally: - finish.set() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - del task - await checkpoint() - gc.collect() - assert ref() is None - - -async def cancellation_unwinds(owners): - await cancellation_case(owners) - - -async def cancellation_during_cleanup(owners): - await cancellation_case(owners, repeated=True) - - -async def cancellation_suppressed(owners): - await cancellation_case(owners, suppress=True) - - -async def registration_and_gc(owners): - original = Value() - original_ref = weakref.ref(original) - registered = owners.prepare(lambda value: value, (original,)) - active = registered.clone_owner() - registered.close() - registered = owners.prepare(lambda: "replacement", ()) - del original - assert active.invoke() is original_ref() - active.close() - gc.collect() - assert original_ref() is None - assert registered.invoke() == "replacement" - registered.close() - - for edge in ("callable", "positional", "keywords"): - - class Callback: - def __call__(self, *args, **kwargs): - pass - - value = Callback() - owner = owners.prepare( - value if edge == "callable" else lambda *a, **k: None, - (value,) if edge == "positional" else (), - {"value": value} if edge == "keywords" else None, - ) - value.owner = owner - value_ref, owner_ref = weakref.ref(value), weakref.ref(owner) - del value, owner - gc.collect() - assert value_ref() is None and owner_ref() is None - assert owners.live == 0 - - finalized = [] - - class Reenter: - def __del__(self): - try: - reentrant.close() - another = owners.prepare(lambda: 42, ()) - finalized.append(another.invoke()) - another.close() - except BaseException as error: - finalized.append(type(error).__name__) - - value = Reenter() - reentrant = owners.prepare(lambda value: None, (value,)) - del value - reentrant.close() - reentrant.close() - assert finalized == [42] - assert owners.live == 0 - - -async def background_and_session(owners): - context = contextvars.ContextVar("background_context", default="initial") - start, finish = asyncio.Event(), asyncio.Event() - payload = {"nested": {"value": "queued"}} - saved = [] - - async def upload(value): - assert context.get() == "submission" - start.set() - await finish.wait() - saved.append(json.dumps(value)) - - session = owners.prepare(lambda value: value, (payload,)) - first_response, second_response = session.clone_owner(), session.clone_owner() - upload_owner = owners.prepare(upload, (payload,), awaited=True) - context.set("submission") - task = asyncio.create_task(upload_owner.invoke()) - context.set("consumer") - upload_owner.close() - first_response.close() - await start.wait() - assert second_response.invoke() is payload - payload["nested"]["value"] = "later" - consumed = json.dumps(payload) - payload["nested"]["after_consumption"] = True - assert "after_consumption" not in consumed - second_response.close() - session.close() - assert owners.live == 0 and not task.done() - finish.set() - await task - assert json.loads(saved[0]) == payload - assert context.get() == "consumer" - - -async def stream_lifecycle(owners): - for terminal in ("exhaustion", "failure", "close"): - nested = {"usage": 0} - item = {"nested": nested} - closed = [] - - async def source(item=item, nested=nested, terminal=terminal, closed=closed): - try: - yield item - nested["usage"] = 12 - if terminal == "failure": - raise ValueError("stream failure") - finally: - closed.append(True) - - stream = source() - pull = owners.prepare(stream.__anext__, (), awaited=True) - yielded = await pull.invoke() - assert yielded is item - shallow, deep = copy.copy(yielded), copy.deepcopy(yielded) - retained = owners.prepare(lambda value: value, (yielded["nested"],)) - yielded["nested"] = {"replacement": True} - if terminal == "close": - close = owners.prepare(stream.aclose, (), awaited=True) - await close.invoke() - await close.invoke() - close.close() - else: - with TestCase().assertRaises(ValueError if terminal == "failure" else StopAsyncIteration): - await pull.invoke() - pull.close() - assert closed == [True] - assert retained.invoke() is nested - assert shallow["nested"] is nested and deep["nested"]["usage"] == 0 - assert nested["usage"] == (0 if terminal == "close" else 12) - retained.close() - - -async def sync_stream_lifecycle(owners): - for terminal in ("exhaustion", "failure", "close"): - value = {"nested": {"usage": 0}} - closed = [] - - def source(value=value, terminal=terminal, closed=closed): - try: - yield value - value["nested"]["usage"] = 12 - if terminal == "failure": - raise ValueError("stream failure") - finally: - closed.append(True) - - stream = source() - pull = owners.prepare(stream.__next__, ()) - assert pull.invoke() is value - saved = owners.prepare(lambda value: value, (value,)) - if terminal == "close": - close = owners.prepare(stream.close, ()) - close.invoke() - close.invoke() - close.close() - else: - with TestCase().assertRaises(ValueError if terminal == "failure" else StopIteration): - pull.invoke() - pull.close() - assert saved.invoke() is value - assert value["nested"]["usage"] == (0 if terminal == "close" else 12) - assert closed == [True] - saved.close() - - -async def repeated_ownership(owners): - refs = [] - for batch in range(8): - gate = asyncio.Event() - - async def work(value, gate=gate): - await gate.wait() - return None - - tasks = [] - for index in range(4): - value = Value() - refs.append(weakref.ref(value)) - owner = owners.prepare(work, (value,), awaited=True) - tasks.append(asyncio.create_task(owner.invoke())) - owner.close() - del value - gate.set() - await asyncio.gather(*tasks) - del tasks - gc.collect() - assert owners.live == 0 - assert all(ref() is None for ref in refs) - - -async def detached_work_after_error(owners): - for raises in (False, True): - entered, release = asyncio.Event(), asyncio.Event() - tasks, observed = [], [] - value = Value() - value.status = "before return" - reference = weakref.ref(value) - - async def consume(argument, entered=entered, release=release, observed=observed): - entered.set() - await release.wait() - observed.append(argument.status) - - def callback(argument, tasks=tasks, consume=consume, raises=raises): - tasks.append(asyncio.create_task(consume(argument))) - if raises: - raise ValueError("after task creation") - - owner = owners.prepare(callback, (value,)) - try: - if raises: - with TestCase().assertRaisesRegex(ValueError, "after task creation"): - owner.invoke() - else: - assert owner.invoke() is None - finally: - owner.close() - del value - try: - await entered.wait() - assert owners.live == 0 and reference() is not None - reference().status = "after invocation" - release.set() - await tasks[0] - assert observed == ["after invocation"] - finally: - release.set() - await asyncio.gather(*tasks, return_exceptions=True) - tasks.clear() - await checkpoint() - gc.collect() - assert reference() is None - - -def run_checked(owners, scenario): - baseline = owners.live - background_failures = [] - - async def run(): - asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context)) - 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 - - 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): - owners = factory if retained else ReferenceFactory() - assert owners.live == 0 - run_checked(owners, globals()[name](owners)) diff --git a/litellm-rust/crates/python-interop/tests/integration/lifecycle.rs b/litellm-rust/crates/python-interop/tests/integration/lifecycle.rs deleted file mode 100644 index e4600e107ff..00000000000 --- a/litellm-rust/crates/python-interop/tests/integration/lifecycle.rs +++ /dev/null @@ -1,96 +0,0 @@ -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::{fixture, rstest}; -use serial_test::serial; - -use crate::support::Backend; -use crate::support::python::run_fixture; -use crate::support::scenarios::{run_scenario_fixture, scenario_scope}; - -#[fixture] -fn component_scope(scenario_scope: Py) -> Py { - Python::attach(|py| { - run_fixture( - py, - scenario_scope.bind(py), - include_str!("../fixtures/callback_components.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/callback_components.py" - ), - ) - .unwrap(); - }); - scenario_scope -} - -#[fixture] -fn integration_scope(scenario_scope: Py) -> Py { - Python::attach(|py| { - run_fixture( - py, - scenario_scope.bind(py), - include_str!("../fixtures/callback_integrations.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/callback_integrations.py" - ), - ) - .unwrap(); - }); - scenario_scope -} - -#[rstest] -#[case::identity_and_ignored_returns("pre_call_identity_and_ignored_returns")] -#[case::mutations_visible_to_later_callbacks("pre_call_mutations_visible_to_later_callbacks")] -#[case::mutation_survives_failure("pre_call_mutation_survives_failure")] -#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] -#[serial(python_interpreter)] -fn pre_call_contract( - component_scope: Py, - #[case] scenario: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> PyResult<()> { - run_scenario_fixture(component_scope, scenario, backend) -} - -#[rstest] -#[case::real_post_call_logging("real_post_call_logging")] -#[case::real_post_call_dict_response("real_post_call_dict_response")] -#[case::real_sync_logging("real_sync_logging")] -#[case::real_sync_logging_hook_failure("real_sync_logging_hook_failure")] -#[case::real_sync_failure_chain("real_sync_failure_chain")] -#[case::real_async_failure_chain("real_async_failure_chain")] -#[case::real_async_logging("real_async_logging")] -#[case::real_copy_boundaries("real_copy_boundaries")] -#[case::real_logging_worker("real_logging_worker")] -#[case::real_sync_stream_copies("real_sync_stream_copies")] -#[case::real_stream_completion("real_stream_completion")] -#[case::real_stream_close("real_stream_close")] -#[case::real_stream_cancellation("real_stream_cancellation")] -#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] -#[serial(python_interpreter)] -fn component_contract( - component_scope: Py, - #[case] scenario: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> PyResult<()> { - run_scenario_fixture(component_scope, scenario, backend) -} - -#[rstest] -#[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_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( - integration_scope: Py, - #[case] scenario: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> PyResult<()> { - run_scenario_fixture(integration_scope, scenario, backend) -} diff --git a/litellm-rust/crates/python-interop/tests/integration/mod.rs b/litellm-rust/crates/python-interop/tests/integration/mod.rs deleted file mode 100644 index 6213338de2d..00000000000 --- a/litellm-rust/crates/python-interop/tests/integration/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -#[path = "../support/mod.rs"] -mod support; - -mod lifecycle; -mod ocr; diff --git a/litellm-rust/crates/python-interop/tests/integration/ocr.rs b/litellm-rust/crates/python-interop/tests/integration/ocr.rs deleted file mode 100644 index 8c38ef61006..00000000000 --- a/litellm-rust/crates/python-interop/tests/integration/ocr.rs +++ /dev/null @@ -1,175 +0,0 @@ -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyTuple}; -use rstest::rstest; -use serial_test::serial; - -use crate::support::python::{InitializedPython, initialized_python, item, scope}; - -fn prepare_pre_call( - py: Python<'_>, - logger: &Bound<'_, PyAny>, - view: &Bound<'_, PyDict>, -) -> PyResult { - let keywords = PyDict::new(py); - keywords.set_item("input", "OCR document processing")?; - keywords.set_item("api_key", py.None())?; - keywords.set_item("additional_args", view)?; - Ok(PreparedCall::new( - InvocationMode::Direct, - logger.getattr("pre_call")?.unbind(), - PyTuple::empty(py).unbind(), - Some(keywords.unbind()), - )) -} - -#[rstest] -#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] -#[serial(python_interpreter)] -fn real_ocr_logging_preserves_execution_roots_and_continues_after_error( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = scope( - py, - c" -from litellm.integrations.custom_logger import CustomLogger - -class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append('retain') - self.view = kwargs['additional_args'] - self.headers = self.view['headers'] - self.body = self.view['complete_input_dict'] - self.snapshot = (self.headers['X-Trace'], self.body['document']['value']) - return {'ignored_replacement': True} - -class MutateThenFail(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append('mutate_then_fail') - view = kwargs['additional_args'] - view['headers']['X-Trace'] = 'mutated' - view['complete_input_dict']['document']['value'] = 'mutated' - view['headers'] = {'X-Trace': 'replacement'} - view['complete_input_dict'] = {'replacement': True} - raise RuntimeError('expected callback failure') - -class Observe(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - order.append('observe') - self.view = kwargs['additional_args'] - self.snapshot = ( - tuple(sorted(self.view['headers'].items())), - self.view['complete_input_dict'].get('replacement'), - 'document' in self.view['complete_input_dict'], - ) - -", - )?; - 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)?; - match invocation.invoke(py)? { - InvocationOutcome::Returned(value) => assert!(value.is_none(py)), - InvocationOutcome::Awaitable(_) => { - panic!("direct binding produced an awaitable outcome") - } - } - drop(invocation); - 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!( - 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" - ); - Ok(()) - }) -} diff --git a/litellm-rust/crates/python-interop/tests/support/callback_owner.rs b/litellm-rust/crates/python-interop/tests/support/callback_owner.rs deleted file mode 100644 index c7f39dad294..00000000000 --- a/litellm-rust/crates/python-interop/tests/support/callback_owner.rs +++ /dev/null @@ -1,114 +0,0 @@ -use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, -}; - -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::*; -use pyo3::pyclass::{PyTraverseError, PyVisit}; -use pyo3::types::{PyDict, PyTuple}; - -#[pyclass] -#[derive(Default)] -pub struct OwnerFactory { - live: Arc, -} - -#[pymethods] -impl OwnerFactory { - #[getter] - fn live(&self) -> usize { - self.live.load(Ordering::SeqCst) - } - - #[pyo3(signature = (callable, positional, keywords=None, awaited=false))] - fn prepare( - &self, - callable: Py, - positional: Py, - keywords: Option>, - awaited: bool, - ) -> Owner { - self.live.fetch_add(1, Ordering::SeqCst); - Owner { - call: Some(PreparedCall::new( - if awaited { - InvocationMode::Await - } else { - InvocationMode::Direct - }, - callable, - positional, - keywords, - )), - live: self.live.clone(), - } - } -} - -#[pyclass(weakref)] -struct Owner { - call: Option, - live: Arc, -} - -impl Owner { - fn release(&mut self) -> Option { - let call = self.call.take(); - if call.is_some() { - self.live.fetch_sub(1, Ordering::SeqCst); - } - call - } -} - -#[pymethods] -impl Owner { - fn invoke(slf: &Bound<'_, Self>) -> PyResult> { - let py = slf.py(); - let call = slf - .borrow() - .call - .as_ref() - .map(|call| call.clone_ref(py)) - .ok_or_else(|| PyRuntimeError::new_err("invocation owner released"))?; - match call.invoke(py)? { - InvocationOutcome::Returned(value) | InvocationOutcome::Awaitable(value) => Ok(value), - } - } - - fn clone_owner(&self, py: Python<'_>) -> PyResult { - let call = self - .call - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("invocation owner released"))?; - self.live.fetch_add(1, Ordering::SeqCst); - Ok(Self { - call: Some(call.clone_ref(py)), - live: self.live.clone(), - }) - } - - fn close(slf: &Bound<'_, Self>) { - let released = slf.borrow_mut().release(); - drop(released); - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Some(call) = &self.call { - call.traverse(visit)?; - } - Ok(()) - } - - fn __clear__(slf: &Bound<'_, Self>) { - Self::close(slf); - } -} - -impl Drop for Owner { - fn drop(&mut self) { - drop(self.release()); - } -} diff --git a/litellm-rust/crates/python-interop/tests/support/mod.rs b/litellm-rust/crates/python-interop/tests/support/mod.rs index 5f39a4df4f8..c0a088c7544 100644 --- a/litellm-rust/crates/python-interop/tests/support/mod.rs +++ b/litellm-rust/crates/python-interop/tests/support/mod.rs @@ -1,9 +1 @@ -mod callback_owner; pub mod python; -pub mod scenarios; - -#[derive(Clone, Copy, Debug)] -pub enum Backend { - Python, - PreparedCall, -} diff --git a/litellm-rust/crates/python-interop/tests/support/scenarios.rs b/litellm-rust/crates/python-interop/tests/support/scenarios.rs deleted file mode 100644 index e59e3ff1873..00000000000 --- a/litellm-rust/crates/python-interop/tests/support/scenarios.rs +++ /dev/null @@ -1,54 +0,0 @@ -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::fixture; - -use super::python::{InitializedPython, initialized_python, run_fixture}; -use super::{Backend, callback_owner}; - -#[fixture] -pub fn scenario_scope(initialized_python: &InitializedPython) -> Py { - initialized_python.attach(|py| { - let globals = PyDict::new(py); - globals - .set_item( - "factory", - Py::new(py, callback_owner::OwnerFactory::default()).unwrap(), - ) - .unwrap(); - globals - .set_item( - "AWAIT_ADAPTER_FILENAME", - litellm_python_interop::AWAIT_ADAPTER_FILENAME - .to_str() - .unwrap(), - ) - .unwrap(); - run_fixture( - py, - &globals, - include_str!("../fixtures/callback_lifecycle.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/callback_lifecycle.py" - ), - ) - .unwrap(); - globals.unbind() - }) -} - -pub fn run_scenario_fixture( - scenario_scope: Py, - scenario: &str, - backend: Backend, -) -> PyResult<()> { - Python::attach(|py| { - let globals = scenario_scope.bind(py); - globals.get_item("run_scenario")?.unwrap().call1(( - scenario, - matches!(backend, Backend::PreparedCall), - globals.get_item("factory")?.unwrap(), - ))?; - Ok(()) - }) -} diff --git a/litellm-rust/crates/python-interop/tests/synthetic/controls.rs b/litellm-rust/crates/python-interop/tests/synthetic/controls.rs deleted file mode 100644 index 56fb36283a5..00000000000 --- a/litellm-rust/crates/python-interop/tests/synthetic/controls.rs +++ /dev/null @@ -1,225 +0,0 @@ -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; - -use crate::support::Backend; -use crate::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>, - backend: Backend, - mode: InvocationMode, - ) -> PyResult { - if matches!(backend, Backend::PreparedCall) { - 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(Backend::Python, Backend::PreparedCall)] backend: Backend, - #[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::()?), - backend, - 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(Backend::Python, Backend::PreparedCall)] backend: Backend, - #[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, - backend, - 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/synthetic/lifecycle.rs b/litellm-rust/crates/python-interop/tests/synthetic/lifecycle.rs deleted file mode 100644 index 39c2ace6bca..00000000000 --- a/litellm-rust/crates/python-interop/tests/synthetic/lifecycle.rs +++ /dev/null @@ -1,411 +0,0 @@ -use std::io::Read; -use std::process::{Command, Stdio}; -use std::time::{Duration, Instant}; - -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; -use rstest::rstest; -use serial_test::{parallel, serial}; - -use crate::support::Backend; -use crate::support::python::{InitializedPython, initialized_python, item, run_fixture}; -use crate::support::scenarios::{run_scenario_fixture, scenario_scope}; - -#[test] -fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> { - let test = concat!( - module_path!(), - "::cold_awaited_adapter_initialization_allows_reentry" - ) - .split_once("::") - .unwrap() - .1; - let child_env = "LITELLM_INTEROP_COLD_REENTRY_CHILD"; - if std::env::var(child_env).as_deref() != Ok(test) { - let mut child = Command::new(std::env::current_exe().unwrap()) - .args(["--exact", test, "--nocapture"]) - .env(child_env, test) - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - let deadline = Instant::now() + Duration::from_secs(15); - loop { - if let Some(status) = child.try_wait().unwrap() { - let mut output = String::new(); - child - .stdout - .take() - .unwrap() - .read_to_string(&mut output) - .unwrap(); - assert!( - status.success(), - "awaited adapter reentry child failed: {status}\n{output}" - ); - assert!( - output.contains("test result: ok. 1 passed; 0 failed; 0 ignored;"), - "awaited adapter reentry child did not run exactly one test:\n{output}" - ); - return Ok(()); - } - if Instant::now() >= deadline { - child.kill().unwrap(); - child.wait().unwrap(); - panic!("awaited adapter reentry did not complete within 15 seconds"); - } - std::thread::sleep(Duration::from_millis(10)); - } - } - - let globals = scenario_scope(&initialized_python()); - Python::attach(|py| { - globals - .bind(py) - .get_item("cold_awaited_adapter_reentry")? - .unwrap() - .call1((globals.bind(py).get_item("factory")?.unwrap(),))?; - Ok(()) - }) -} - -#[rstest] -#[case::awaitable_kinds("awaitable_kinds")] -#[case::identity_and_context("identity_and_context")] -#[case::exceptions("exceptions")] -#[case::exception_ownership("exception_ownership")] -#[case::cancellation_before_start("cancellation_before_start")] -#[case::cancellation_unwinds("cancellation_unwinds")] -#[case::cancellation_during_cleanup("cancellation_during_cleanup")] -#[case::cancellation_suppressed("cancellation_suppressed")] -#[case::registration_and_gc("registration_and_gc")] -#[case::background_and_session("background_and_session")] -#[case::stream_lifecycle("stream_lifecycle")] -#[case::sync_stream_lifecycle("sync_stream_lifecycle")] -#[case::repeated_ownership("repeated_ownership")] -#[case::detached_work_after_error("detached_work_after_error")] -#[serial(python_interpreter)] -fn lifecycle_contract( - scenario_scope: Py, - #[case] scenario: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> PyResult<()> { - run_scenario_fixture(scenario_scope, scenario, backend) -} - -#[rstest] -#[case::retained_lifetime("deferred_lifetime", "identity")] -#[case::prepared_ownership("deferred_lifetime", "missing_handoff")] -#[case::externally_owned_retained("borrowed_lifetime", "identity")] -#[case::original_coroutine("direct_coroutine", "identity")] -#[case::passthrough_coroutine("direct_coroutine", "result_passthrough")] -#[serial(python_interpreter)] -fn control_contract( - scenario_scope: Py, - #[case] witness: &str, - #[case] control: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, - #[values(false, true)] awaited: bool, -) -> PyResult<()> { - run_control_fixture(scenario_scope, witness, control, backend, awaited) -} - -#[rstest] -#[case::expired_borrow("deferred_lifetime")] -#[case::externally_owned_borrow("borrowed_lifetime")] -#[serial(python_interpreter)] -fn weak_control( - scenario_scope: Py, - #[case] witness: &str, - #[values(false, true)] awaited: bool, -) -> PyResult<()> { - run_control_fixture(scenario_scope, witness, "weak", Backend::Python, awaited) -} - -#[rstest] -#[case::retained("identity")] -#[case::missing_handoff("missing_handoff")] -#[serial(python_interpreter)] -fn pending_handoff_control( - scenario_scope: Py, - #[case] control: &str, - #[values(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> PyResult<()> { - run_control_fixture(scenario_scope, "pending_handoff", control, backend, true) -} - -fn run_control_fixture( - scenario_scope: Py, - witness: &str, - control: &str, - backend: Backend, - awaited: bool, -) -> PyResult<()> { - Python::attach(|py| { - let globals = scenario_scope.bind(py); - run_fixture( - py, - globals, - include_str!("../fixtures/callback_controls.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/callback_controls.py" - ), - )?; - globals.get_item("run_control")?.unwrap().call1(( - witness, - control, - matches!(backend, Backend::PreparedCall), - awaited, - globals.get_item("factory")?.unwrap(), - ))?; - Ok(()) - }) -} - -fn invoke_direct_callback( - py: Python<'_>, - globals: &Bound<'_, PyDict>, - callback: &str, - argument: &str, - backend: Backend, -) -> PyResult> { - let args = PyTuple::new(py, [item(globals, argument)])?; - if matches!(backend, Backend::Python) { - 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(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> 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", backend)?.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(Backend::Python, Backend::PreparedCall)] backend: Backend, -) -> 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", backend)?.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<()> { - use litellm_python_interop::{InvocationMode, PreparedCall}; - use pyo3::types::PyTuple; - let _ = initialized_python; - let (call, reference) = Python::attach(|py| -> PyResult<_> { - let globals = PyDict::new(py); - py.run( - c"import weakref\nclass Value: pass\nvalue = Value()\nreference = weakref.ref(value)", - Some(&globals), - None, - )?; - let value = globals.get_item("value")?.unwrap(); - let call = PreparedCall::new( - InvocationMode::Direct, - py.None(), - PyTuple::new(py, [value])?.unbind(), - None, - ); - let reference = globals.get_item("reference")?.unwrap().unbind(); - globals.del_item("value")?; - Ok((call, reference)) - })?; - drop(call); - Python::attach(|py| { - assert!(reference.call0(py)?.is_none(py)); - Ok(()) - }) -} - -#[rstest] -#[case::direct(false)] -#[case::awaited(true)] -#[parallel(python_interpreter)] -fn outcome_identifies_binding(scenario_scope: Py, #[case] awaited: bool) -> PyResult<()> { - use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; - use pyo3::types::PyTuple; - Python::attach(|py| { - let callback = py.eval(c"lambda: None", Some(scenario_scope.bind(py)), None)?; - let call = PreparedCall::new( - if awaited { - InvocationMode::Await - } else { - InvocationMode::Direct - }, - callback.unbind(), - PyTuple::empty(py).unbind(), - None, - ); - match call.invoke(py)? { - InvocationOutcome::Returned(value) => { - assert!(!awaited); - assert!(value.is_none(py)); - } - InvocationOutcome::Awaitable(value) => { - assert!(awaited); - value.call_method0(py, "close")?; - } - } - Ok(()) - }) -} - -#[rstest] -#[parallel(python_interpreter)] -fn awaited_raise_surfaces_when_driven(scenario_scope: Py) -> PyResult<()> { - use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; - use pyo3::types::PyTuple; - Python::attach(|py| { - let globals = scenario_scope.bind(py); - py.run( - c"events = []\nerror = ValueError('await failure')\nasync def callback():\n events.append('started')\n raise error\ndef drive(coroutine):\n try:\n coroutine.send(None)\n return False\n except BaseException as caught:\n return caught is error\n", - Some(globals), - None, - )?; - let started = || -> PyResult { globals.get_item("events")?.unwrap().len() }; - let call = PreparedCall::new( - InvocationMode::Await, - globals.get_item("callback")?.unwrap().unbind(), - PyTuple::empty(py).unbind(), - None, - ); - let pending = match call.invoke(py)? { - InvocationOutcome::Awaitable(value) => value, - InvocationOutcome::Returned(_) => panic!("await binding produced a settled outcome"), - }; - assert_eq!(started()?, 0); - assert!( - globals - .get_item("drive")? - .unwrap() - .call1((pending,))? - .extract::()? - ); - assert_eq!(started()?, 1); - Ok(()) - }) -} diff --git a/litellm-rust/crates/python-interop/tests/synthetic/mod.rs b/litellm-rust/crates/python-interop/tests/synthetic/mod.rs index 2e4f9bc28ef..053bb5f2883 100644 --- a/litellm-rust/crates/python-interop/tests/synthetic/mod.rs +++ b/litellm-rust/crates/python-interop/tests/synthetic/mod.rs @@ -1,8 +1,4 @@ #[path = "../support/mod.rs"] mod support; -mod controls; -mod lifecycle; -mod patterns; -mod prepared_call; mod primitives; diff --git a/litellm-rust/crates/python-interop/tests/synthetic/patterns.rs b/litellm-rust/crates/python-interop/tests/synthetic/patterns.rs deleted file mode 100644 index f9489af6e94..00000000000 --- a/litellm-rust/crates/python-interop/tests/synthetic/patterns.rs +++ /dev/null @@ -1,712 +0,0 @@ -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; - -use crate::support::Backend; -use crate::support::python::{InitializedPython, initialized_python, item, run_fixture, scope}; - -#[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/synthetic/prepared_call.rs b/litellm-rust/crates/python-interop/tests/synthetic/prepared_call.rs deleted file mode 100644 index 6ff35c344c0..00000000000 --- a/litellm-rust/crates/python-interop/tests/synthetic/prepared_call.rs +++ /dev/null @@ -1,350 +0,0 @@ -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; -use pyo3::exceptions::{PyAssertionError, PyKeyboardInterrupt, PyValueError}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyTuple}; -use rstest::rstest; -use serial_test::parallel; - -use crate::support::python::{InitializedPython, initialized_python, item, run_fixture, scope}; - -#[rstest] -#[parallel(python_interpreter)] -fn retains_aliases_mutations_and_original_result( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = scope( - py, - c" -def callback(data, *, alias): - observed.append(data['nested'] is alias) - saved.append(data) - alias['value'] = 'during' - return data -", - )?; - 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", &shared)?; - let invocation = PreparedCall::new( - InvocationMode::Direct, - item(&globals, "callback").unbind(), - PyTuple::new(py, [&payload])?.unbind(), - Some(keywords.unbind()), - ); - let result = invoke_direct(&invocation, py)?; - assert!(result.bind(py).is(&payload)); - drop(invocation); - 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(()) - }) -} - -#[rstest] -#[parallel(python_interpreter)] -fn preserves_exception_identity_cause_traceback_and_prior_mutation( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = scope( - py, - c" -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, [&payload])?.unbind(), - None, - ); - let error = invoke_direct(&invocation, py).unwrap_err(); - assert!(error.value(py).is(original.value(py))); - drop(invocation); - 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(()) - }) -} - -#[rstest] -#[parallel(python_interpreter)] -fn returns_coroutine_without_executing_it(initialized_python: &InitializedPython) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = scope( - py, - c" -async def work(): - started.append(True) -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(), - PyTuple::empty(py).unbind(), - None, - ); - let result = invoke_direct(&invocation, py)?; - 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(()) - }) -} - -#[pyfunction] -fn reenter(py: Python<'_>, callback: Py, payload: Py) -> PyResult> { - invoke_direct( - &PreparedCall::new( - InvocationMode::Direct, - callback, - PyTuple::new(py, [payload])?.unbind(), - None, - ), - py, - ) -} - -#[rstest] -#[parallel(python_interpreter)] -fn preserves_current_context_thread_and_reentry( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = scope( - py, - c" -import threading -def inner(data): - observed.append((context.get(), threading.get_ident())) - data['inner'] = True - context.set('inner') - return data -def outer(): - 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, - item(&globals, "outer").unbind(), - PyTuple::empty(py).unbind(), - 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(()) - }) -} - -#[rstest] -#[parallel(python_interpreter)] -fn owns_arguments_until_release_and_preserves_callback_retention( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - let (invocation, globals) = Python::attach(|py| { - let globals = scope( - py, - c" -class Value: - pass -class Callback: - def __call__(self, value, *, other): - saved.append(value) - 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", other)?; - let invocation = PreparedCall::new( - InvocationMode::Direct, - callback.unbind(), - PyTuple::new(py, [value])?.unbind(), - Some(keywords.unbind()), - ); - Ok::<_, PyErr>((invocation, globals.unbind())) - })?; - Python::attach(|py| { - let globals = globals.bind(py); - 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); - 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(()) - }) -} - -#[rstest] -#[parallel(python_interpreter)] -fn checked_runner_rejects_unhandled_background_failures( - initialized_python: &InitializedPython, -) -> PyResult<()> { - let _ = initialized_python; - Python::attach(|py| { - let globals = PyDict::new(py); - run_fixture( - py, - &globals, - include_str!("../fixtures/callback_lifecycle.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/callback_lifecycle.py" - ), - )?; - py.run( - c" -async def fail(): - raise RuntimeError('background task regression') - -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(()) - }) -} - -fn invoke_direct(call: &PreparedCall, py: Python<'_>) -> PyResult> { - match call.invoke(py)? { - InvocationOutcome::Returned(value) => Ok(value), - InvocationOutcome::Awaitable(_) => panic!("direct binding produced an awaitable outcome"), - } -} diff --git a/litellm-rust/crates/python-interop/tests/synthetic/primitives.rs b/litellm-rust/crates/python-interop/tests/synthetic/primitives.rs index 5b948b6246e..e6ceda6051f 100644 --- a/litellm-rust/crates/python-interop/tests/synthetic/primitives.rs +++ b/litellm-rust/crates/python-interop/tests/synthetic/primitives.rs @@ -2,7 +2,7 @@ use rstest::rstest; use serde_json::{Value, json}; use serial_test::parallel; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; +use litellm_python_interop::{from_py, to_py}; use crate::support::python::{InitializedPython, initialized_python}; @@ -18,13 +18,3 @@ fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &I assert_eq!(actual, expected); }); } - -#[rstest] -#[parallel(python_interpreter)] -fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) { - let before = release_count(); - let result = python.attach(|py| release_gil(py, || 42)); - - assert_eq!(result, 42); - assert_eq!(release_count(), before + 1); -}