From 0750ef81ca3672a77133e11f6e91486083306adf Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 19:05:47 -0700 Subject: [PATCH] refactor(rust): unify call lifecycle execution --- litellm-rust/Cargo.lock | 24 + litellm-rust/Cargo.toml | 4 + .../src/audio_transcription/hooks.rs | 192 +++---- .../ai-gateway/src/audio_transcription/mod.rs | 127 ++++- .../src/audio_transcription/prepare.rs | 7 +- .../src/routes/responses/service.rs | 33 +- litellm-rust/crates/core/Cargo.toml | 2 + .../core/src/audio_transcription/handler.rs | 37 +- .../core/src/audio_transcription/lifecycle.rs | 69 +++ .../core/src/audio_transcription/mod.rs | 7 +- .../core/src/caching/in_memory_cache.rs | 258 --------- litellm-rust/crates/core/src/caching/mod.rs | 1 - .../crates/core/src/call_lifecycle/cache.rs | 71 +++ .../core/src/call_lifecycle/execution.rs | 220 ++++++++ .../crates/core/src/call_lifecycle/host.rs | 42 +- .../crates/core/src/call_lifecycle/mod.rs | 417 +------------- .../core/src/call_lifecycle/provider.rs | 516 ++++++++++++++++++ .../crates/core/src/call_lifecycle/types.rs | 35 +- .../core/src/call_lifecycle/workflow.rs | 250 +++++++++ .../core/src/chat_completions/handler.rs | 29 +- .../core/src/chat_completions/lifecycle.rs | 69 +++ .../crates/core/src/chat_completions/mod.rs | 9 +- litellm-rust/crates/core/src/lib.rs | 1 - .../crates/core/src/messages/handler.rs | 29 +- .../crates/core/src/messages/lifecycle.rs | 62 +++ litellm-rust/crates/core/src/messages/mod.rs | 5 +- litellm-rust/crates/core/src/ocr/handler.rs | 55 +- litellm-rust/crates/core/src/ocr/hooks.rs | 89 +-- litellm-rust/crates/core/src/ocr/lifecycle.rs | 479 +++++----------- litellm-rust/crates/core/src/ocr/mod.rs | 3 +- litellm-rust/crates/core/src/ocr/wire.rs | 7 +- .../core/src/providers/bedrock/aws_base.rs | 21 +- .../core/src/responses/instrumentation.rs | 83 +-- .../crates/core/tests/host_lifecycle.rs | 60 ++ litellm-rust/crates/core/tests/ocr.rs | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 1 + .../python-bridge/src/lifecycle/bindings.rs | 16 +- .../python-bridge/src/lifecycle/completed.rs | 281 ++++++++++ .../python-bridge/src/lifecycle/contract.rs | 163 ++++++ .../python-bridge/src/lifecycle/dispatch.rs | 9 +- .../python-bridge/src/lifecycle/handle.rs | 3 +- .../crates/python-bridge/src/lifecycle/mod.rs | 3 + .../python-bridge/src/lifecycle/request.rs | 63 +++ .../python-bridge/src/lifecycle/runner.rs | 43 +- .../python-bridge/src/lifecycle/state.rs | 23 +- .../python-bridge/src/lifecycle/tests.rs | 17 +- .../src/routes/chat_completions/lifecycle.rs | 59 +- .../src/routes/chat_completions/value.rs | 2 +- .../src/routes/messages/lifecycle.rs | 51 +- .../python-bridge/src/routes/ocr/callbacks.rs | 5 +- .../python-bridge/src/routes/ocr/lifecycle.rs | 17 +- .../python-bridge/src/routes/ocr/value.rs | 8 +- .../src/routes/transcription/lifecycle.rs | 54 +- litellm/rust_bridge/_native.pyi | 106 +++- litellm/rust_bridge/catalog.py | 2 +- .../rust_bridge/test_chat_completions.py | 12 +- .../test_route_foundation.py | 18 +- 58 files changed, 2791 insertions(+), 1482 deletions(-) create mode 100644 litellm-rust/crates/core/src/audio_transcription/lifecycle.rs delete mode 100644 litellm-rust/crates/core/src/caching/in_memory_cache.rs delete mode 100644 litellm-rust/crates/core/src/caching/mod.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/cache.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/execution.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/provider.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/workflow.rs create mode 100644 litellm-rust/crates/core/src/chat_completions/lifecycle.rs create mode 100644 litellm-rust/crates/core/src/messages/lifecycle.rs create mode 100644 litellm-rust/crates/python-bridge/src/lifecycle/completed.rs create mode 100644 litellm-rust/crates/python-bridge/src/lifecycle/contract.rs create mode 100644 litellm-rust/crates/python-bridge/src/lifecycle/request.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7e3d25e9c5d..fe93c3d094a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1925,6 +1925,26 @@ dependencies = [ "tracing", ] +[[package]] +name = "litellm-cache" +version = "0.1.0" +dependencies = [ + "rstest", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.19", +] + +[[package]] +name = "litellm-cache-memory" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "rstest", + "serde_json", +] + [[package]] name = "litellm-config" version = "0.1.0" @@ -1952,6 +1972,8 @@ dependencies = [ "data-url", "futures-util", "gcp_auth", + "litellm-cache", + "litellm-cache-memory", "mime_guess", "moka", "rand 0.8.7", @@ -1980,6 +2002,8 @@ version = "0.1.0" dependencies = [ "criterion", "futures-util", + "litellm-cache", + "litellm-cache-memory", "litellm-core", "litellm-python-interop", "litellm-token-counter", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 5c72c86d6ef..3b3e3454df3 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,5 +1,7 @@ [workspace] members = [ + "crates/cache", + "crates/cache-memory", "crates/core", "crates/token-counter", "crates/config", @@ -20,6 +22,8 @@ bytes = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } +litellm-cache = { path = "crates/cache" } +litellm-cache-memory = { path = "crates/cache-memory" } litellm-token-counter = { path = "crates/token-counter" } litellm-config = { path = "crates/config" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index b17f17de11f..1249373bc4c 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -1,12 +1,10 @@ -use litellm_core::audio_transcription::{ - AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, - prepare_audio_transcription_provider_call, -}; -use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use litellm_core::error::Error; use serde_json::{Map, Value, json}; -use std::future::Future; -use std::pin::Pin; + +use litellm_core::call_lifecycle::provider::{ + ProviderHookFuture, ProviderHooks, ProviderRequest, ProviderResponse, +}; +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; +use litellm_core::error::Error; use super::types::PreparedAudioTranscriptionRequest; use crate::integrations::custom_guardrail::{ @@ -23,25 +21,25 @@ pub(crate) struct AudioTranscriptionLifecycleHooks { logger_runner: CustomLoggerRunner, guardrail_runner: CustomGuardrailRunner, request_metadata: RequestMetadata, + provider: String, } -type AudioFuture<'a, T> = Pin> + Send + 'a>>; -type AudioLogFuture<'a> = Pin + Send + 'a>>; - impl AudioTranscriptionLifecycleHooks { pub(crate) fn new( logger_runner: CustomLoggerRunner, guardrail_runner: CustomGuardrailRunner, request_metadata: RequestMetadata, + provider: String, ) -> Self { Self { logger_runner, guardrail_runner, request_metadata, + provider, } } - async fn run_pre_call_guardrails( + pub(crate) async fn run_pre_call_guardrails( &self, request: PreparedAudioTranscriptionRequest, ) -> Result { @@ -85,39 +83,10 @@ impl AudioTranscriptionLifecycleHooks { }) } - async fn prepare_provider_request( - &self, - request: PreparedAudioTranscriptionRequest, - ) -> Result { - let PreparedAudioTranscriptionRequest { - model, - custom_llm_provider, - audio, - api_key, - api_base, - extra_headers, - optional_params, - timeout, - .. - } = request; - let provider_request = - prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: Some(&custom_llm_provider), - extra_headers, - optional_params, - timeout, - })?; - self.run_during_call_guardrails(provider_request).await - } - async fn run_during_call_guardrails( &self, - request: ProviderAudioTranscriptionRequest, - ) -> Result { + request: ProviderRequest, + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -126,10 +95,10 @@ impl AudioTranscriptionLifecycleHooks { .run_during_call( &guardrail_context(&self.request_metadata), GuardrailRequest::new(json!({ - "model": request.model(), - "custom_llm_provider": request.custom_llm_provider(), - "url": request.url(), - "body": request.body(), + "model": &request.model, + "custom_llm_provider": &self.provider, + "url": &request.url, + "body": &request.body, })), ) .await @@ -142,7 +111,7 @@ impl AudioTranscriptionLifecycleHooks { let body = data.remove("body").ok_or_else(|| { Error::InvalidRequest("audio transcription guardrail removed body".to_string()) })?; - Ok(request.with_body(body)) + Ok(ProviderRequest { body, ..request }) } fn logging_payload( @@ -172,82 +141,65 @@ impl AudioTranscriptionLifecycleHooks { messages: None, } } + pub(crate) async fn log_success( + &self, + context: &CallLifecycleContext, + response: &Value, + timing: &CallLifecycleTiming, + ) { + if self.logger_runner.is_empty() { + return; + } + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ), + &CallbackValue::new("audio_transcription", response.clone()), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + } + + pub(crate) async fn log_failure( + &self, + context: &CallLifecycleContext, + error: &Error, + timing: &CallLifecycleTiming, + ) { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ) + .with_failure_error(logging_error.clone()), + Some(&CallbackValue::new( + "error", + json!({"message": logging_error.message, "kind": logging_error.kind}), + )), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + } } -impl CallLifecycleHooks - for AudioTranscriptionLifecycleHooks -{ - type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>; - type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>; - type SuccessFuture<'a> = AudioLogFuture<'a>; - type FailureFuture<'a> = AudioLogFuture<'a>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: PreparedAudioTranscriptionRequest, - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { self.run_pre_call_guardrails(request).await }) +impl ProviderHooks for AudioTranscriptionLifecycleHooks { + fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> { + Box::pin(async move { self.run_during_call_guardrails(request).await }) } - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: PreparedAudioTranscriptionRequest, - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { self.prepare_provider_request(request).await }) - } - - fn async_log_success_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - response: &'a Value, - timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - Box::pin(async move { - if self.logger_runner.is_empty() { - return; - } - self.logger_runner - .async_log_success_event( - &ModelCallDetails::from_standard_logging_payload( - self.logging_payload(context, timing), - ), - &CallbackValue::new("audio_transcription", response.clone()), - CallbackTiming::new(timing.start_time, timing.end_time), - ) - .await; - }) - } - - fn async_log_failure_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - error: &'a Error, - timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - Box::pin(async move { - if self.logger_runner.is_empty() { - return; - } - let logging_error = LoggingError { - message: error.to_string(), - kind: core_error_kind(error).to_string(), - }; - self.logger_runner - .async_log_failure_event( - &ModelCallDetails::from_standard_logging_payload( - self.logging_payload(context, timing), - ) - .with_failure_error(logging_error.clone()), - Some(&CallbackValue::new( - "error", - json!({"message": logging_error.message, "kind": logging_error.kind}), - )), - CallbackTiming::new(timing.start_time, timing.end_time), - ) - .await; - }) + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse> { + Box::pin(async move { Ok(response) }) } } diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index 03d621b8414..0fd18b28edc 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -1,22 +1,137 @@ -use litellm_core::Error; -use litellm_core::audio_transcription::execute_audio_transcription_provider_call; -use litellm_core::call_lifecycle::CallLifecycle; +use std::sync::Mutex; +use std::time::{SystemTime, UNIX_EPOCH}; + use serde_json::Value; +use litellm_core::Error; +use litellm_core::audio_transcription::lifecycle::{ + AudioTranscriptionCall, OwnedAudioTranscriptionRequest, +}; +use litellm_core::call_lifecycle::CallLifecycleRequest; +use litellm_core::call_lifecycle::host::{LifecycleBackend, LifecycleBackendFuture, drive}; +use litellm_core::call_lifecycle::provider::{ + CompletedOperation, CompletedReply, CompletedWorkflow, ProviderHooks, ProviderOptions, +}; +use litellm_core::call_lifecycle::workflow::LifecycleOperation; + mod hooks; mod prepare; mod types; pub use types::AudioTranscriptionRequest; +use hooks::AudioTranscriptionLifecycleHooks; use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call}; +use types::PreparedAudioTranscriptionRequest; pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { let PreparedAudioTranscriptionCall { request, hooks } = prepare_audio_transcription_call(request); - CallLifecycle::default() - .run_request(request, &hooks, execute_audio_transcription_provider_call) - .await + let mut call = AudioTranscriptionCall::new(CompletedWorkflow::default(), false); + drive( + &mut call, + &AudioTranscriptionBackend { + request: Mutex::new(Some(request)), + hooks, + }, + ) + .await +} + +struct AudioTranscriptionBackend { + request: Mutex>, + hooks: AudioTranscriptionLifecycleHooks, +} + +impl LifecycleBackend, CompletedReply> + for AudioTranscriptionBackend +{ + fn invoke( + &self, + operation: CompletedOperation, + ) -> LifecycleBackendFuture<'_, CompletedReply> { + Box::pin(async move { + match operation { + CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => { + let request = self + .request + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .ok_or_else(|| { + Error::InvalidRequest("request was already projected".into()) + }); + CompletedReply::Request(match request { + Ok(request) => { + let context = request.lifecycle_context(); + let started = epoch_seconds(); + match self.hooks.run_pre_call_guardrails(request).await { + Ok(request) => Ok(owned_request(request)), + Err(error) => { + let timing = + litellm_core::call_lifecycle::CallLifecycleTiming::new( + started, + epoch_seconds(), + ); + self.hooks.log_failure(&context, &error, &timing).await; + Err(error) + } + } + } + Err(error) => Err(error), + }) + } + CompletedOperation::Lifecycle(LifecycleOperation::Success { + context, + response, + timing, + }) => { + self.hooks + .log_success(&context, response.as_ref(), &timing) + .await; + CompletedReply::Lifecycle(Ok(())) + } + CompletedOperation::Lifecycle(LifecycleOperation::Failure { + context, + error, + timing, + }) => { + self.hooks.log_failure(&context, &error, &timing).await; + CompletedReply::Lifecycle(Ok(())) + } + CompletedOperation::Lifecycle(_) => CompletedReply::Lifecycle(Ok(())), + CompletedOperation::BeforeRequest(request) => { + CompletedReply::BeforeRequest(self.hooks.before_request(request).await) + } + CompletedOperation::AfterResponse(response) => { + CompletedReply::AfterResponse(self.hooks.after_response(response).await) + } + } + }) + } +} + +fn owned_request(request: PreparedAudioTranscriptionRequest) -> OwnedAudioTranscriptionRequest { + OwnedAudioTranscriptionRequest { + options: ProviderOptions { + model: request.model, + litellm_call_id: Some(request.litellm_call_id), + api_key: request.api_key, + api_base: request.api_base, + custom_llm_provider: Some(request.custom_llm_provider), + extra_headers: request.extra_headers, + timeout: request.timeout, + }, + audio: request.audio, + optional_params: request.optional_params, + } +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) } #[cfg(test)] diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs index a475d58635f..b184c9ee533 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs @@ -25,10 +25,12 @@ pub(crate) fn prepare_audio_transcription_call( model: request.model, custom_llm_provider: "bedrock", }); + let model = provider_info.model.to_string(); + let provider = provider_info.custom_llm_provider.to_string(); PreparedAudioTranscriptionCall { request: PreparedAudioTranscriptionRequest { - model: provider_info.model.to_string(), - custom_llm_provider: provider_info.custom_llm_provider.to_string(), + model, + custom_llm_provider: provider.clone(), litellm_call_id: call_id, audio: request.audio, api_key: request.api_key.map(str::to_string), @@ -41,6 +43,7 @@ pub(crate) fn prepare_audio_transcription_call( CustomLoggerRunner::new(request.callbacks), CustomGuardrailRunner::new(request.guardrails), request.request_metadata, + provider, ), } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs index e8f840c0c8e..50d21f3eda8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs @@ -3,7 +3,6 @@ use std::time::Duration; use futures_util::{Sink, Stream}; use litellm_core::Error; -use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext}; use litellm_core::responses::instrumentation::{ ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome, ResponsesWsMetadata, @@ -55,24 +54,20 @@ where }, )); let observer_instrumentation = Arc::clone(&instrumentation); - let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id); - let result = CallLifecycle::default() - .run(context, (), instrumentation.as_ref(), |_| async move { - crate::io::responses_ws::async_responses_websocket( - provider_model, - params.api_key.as_deref(), - params.api_base.as_deref(), - first_frame, - idle_timeout, - move |event| { - observer_instrumentation.observe(event); - }, - client_in, - client_out, - ) - .await - }) - .await; + let result = crate::io::responses_ws::async_responses_websocket( + provider_model, + params.api_key.as_deref(), + params.api_base.as_deref(), + first_frame, + idle_timeout, + move |event| { + observer_instrumentation.observe(event); + }, + client_in, + client_out, + ) + .await; + instrumentation.record_outcome(result.is_ok()); let outcome = instrumentation.take_or_build_outcome(result.is_ok()); dispatch_outcome(loggers, outcome).await; result diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 09c526f73cf..3b4dfd87665 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -8,6 +8,8 @@ autotests = false [dependencies] bytes.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true futures-util.workspace = true base64.workspace = true azure_core.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 9a96b9d1140..318ca2fae0c 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,5 +1,8 @@ use serde_json::Value; +use crate::call_lifecycle::provider::{ + NoopProviderHooks, ProviderHooks, ProviderRequest, ProviderResponse, +}; use crate::error::Error; use crate::http_utils::{http_request, truncate_error_body}; @@ -10,6 +13,28 @@ use super::types::ProviderAudioTranscriptionRequest; pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { + execute_with_hooks(request, &NoopProviderHooks).await +} + +pub(super) async fn execute_with_hooks( + request: ProviderAudioTranscriptionRequest, + hooks: &dyn ProviderHooks, +) -> Result { + let changed = hooks + .before_request(ProviderRequest { + model: request.model.clone(), + url: request.url.clone(), + headers: request.upstream_headers.clone(), + body: request.body.clone(), + }) + .await?; + let request = ProviderAudioTranscriptionRequest { + model: changed.model, + url: changed.url, + upstream_headers: changed.headers, + body: changed.body, + ..request + }; let body = serde_json::to_vec(&request.body) .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; let headers = signed_headers(&request, &body).await?; @@ -28,9 +53,17 @@ pub async fn execute_audio_transcription_provider_call( .text() .await .map_err(|error| Error::Network(error.to_string()))?; - if !status.is_success() { - return Err(Error::Http { + let observed = hooks + .after_response(ProviderResponse { status: status.as_u16(), + body: text, + }) + .await?; + let observed_status = observed.status; + let text = observed.body; + if !(200..300).contains(&observed_status) { + return Err(Error::Http { + status: observed_status, body: truncate_error_body(&text), }); } diff --git a/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs b/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs new file mode 100644 index 00000000000..6e744054dd5 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use serde_json::Value; + +use super::types::AudioTranscriptionRequest; +use crate::call_lifecycle::provider::{ + CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions, +}; +use crate::call_lifecycle::workflow::WorkflowFuture; + +pub struct OwnedAudioTranscriptionRequest { + pub options: ProviderOptions, + pub audio: Value, + pub optional_params: serde_json::Map, +} + +impl From> for OwnedAudioTranscriptionRequest { + fn from(request: AudioTranscriptionRequest<'_>) -> Self { + Self { + options: ProviderOptions { + model: request.model.to_owned(), + litellm_call_id: None, + api_key: request.api_key.map(str::to_owned), + api_base: request.api_base.map(str::to_owned), + custom_llm_provider: request.custom_llm_provider.map(str::to_owned), + extra_headers: request.extra_headers, + timeout: request.timeout, + }, + audio: request.audio, + optional_params: request.optional_params, + } + } +} + +pub struct AudioTranscriptionRoute; +pub type AudioTranscriptionCall = CompletedCall; + +impl CompletedRoute for AudioTranscriptionRoute { + type Request = OwnedAudioTranscriptionRequest; + type Response = Value; + + fn run( + request: Self::Request, + hooks: Arc, + ) -> WorkflowFuture { + Box::pin(async move { + let options = request.options; + let request = AudioTranscriptionRequest { + model: &options.model, + api_key: options.api_key.as_deref(), + api_base: options.api_base.as_deref(), + custom_llm_provider: options.custom_llm_provider.as_deref(), + extra_headers: options.extra_headers, + timeout: options.timeout, + audio: request.audio, + optional_params: request.optional_params, + }; + super::handler::execute_with_hooks( + super::prepare::prepare_audio_transcription_provider_call(request)?, + hooks.as_ref(), + ) + .await + }) + } + + fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext { + request.options.lifecycle_context("audio_transcription") + } +} diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 84042a3629d..eb4335c1d96 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -1,6 +1,7 @@ use crate::Error; mod client; mod handler; +pub mod lifecycle; mod prepare; pub mod transformation; pub mod types; @@ -13,8 +14,10 @@ pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { - execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) - .await + crate::call_lifecycle::provider::run_completed::( + request.into(), + ) + .await } pub fn admit( diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs deleted file mode 100644 index 45d4bd69b79..00000000000 --- a/litellm-rust/crates/core/src/caching/in_memory_cache.rs +++ /dev/null @@ -1,258 +0,0 @@ -use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashMap}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; -const DEFAULT_TTL: Duration = Duration::from_secs(600); - -pub struct InMemoryCache { - pub cache_dict: HashMap, - pub ttl_dict: HashMap, - pub expiration_heap: BinaryHeap>, - pub max_size_in_memory: usize, - pub default_ttl: Duration, - now: Box Duration + Send + Sync>, -} - -impl Default for InMemoryCache { - fn default() -> Self { - Self::new(None, None) - } -} - -impl InMemoryCache { - pub fn new(max_size_in_memory: Option, default_ttl: Option) -> Self { - Self::with_clock(max_size_in_memory, default_ttl, || { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - }) - } - - pub fn with_clock( - max_size_in_memory: Option, - default_ttl: Option, - now: impl Fn() -> Duration + Send + Sync + 'static, - ) -> Self { - Self { - cache_dict: HashMap::new(), - ttl_dict: HashMap::new(), - expiration_heap: BinaryHeap::new(), - max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), - default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - now: Box::new(now), - } - } - - pub fn evict_cache(&mut self) { - if self.max_size_in_memory == 0 { - return; - } - - let current_time = (self.now)(); - while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() { - if self.ttl_dict.get(&key).copied() != Some(expiration_time) { - self.expiration_heap.pop(); - } else if expiration_time <= current_time { - self.expiration_heap.pop(); - self.remove_key(&key); - } else { - break; - } - } - - while self.cache_dict.len() >= self.max_size_in_memory { - let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else { - break; - }; - if self.ttl_dict.get(&key).copied() == Some(expiration_time) { - self.remove_key(&key); - } - } - } - - pub fn allow_ttl_override(&self, key: &str) -> bool { - match self.ttl_dict.get(key).copied() { - None => true, - Some(expiration_time) => expiration_time < (self.now)(), - } - } - - pub fn set_cache(&mut self, key: impl Into, value: V, ttl: Option) { - if self.max_size_in_memory == 0 { - return; - } - - self.evict_cache(); - let key = key.into(); - self.cache_dict.insert(key.clone(), value); - if self.allow_ttl_override(&key) { - let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl); - self.ttl_dict.insert(key.clone(), expiration_time); - self.expiration_heap.push(Reverse((expiration_time, key))); - } - } - - // Generic values intentionally omit Python's per-item size check. - pub fn get_cache(&mut self, key: &str) -> Option { - if self.cache_dict.contains_key(key) { - if self.is_key_expired(key) { - self.remove_key(key); - return None; - } - return self.cache_dict.get(key).cloned(); - } - None - } - - pub fn get_ttl(&self, key: &str) -> Option { - self.ttl_dict.get(key).copied() - } - - pub fn delete_cache(&mut self, key: &str) { - self.remove_key(key); - } - - pub fn flush_cache(&mut self) { - self.cache_dict.clear(); - self.ttl_dict.clear(); - self.expiration_heap.clear(); - } - - fn is_key_expired(&self, key: &str) -> bool { - self.ttl_dict - .get(key) - .is_some_and(|expiration_time| *expiration_time < (self.now)()) - } - - fn remove_key(&mut self, key: &str) { - self.cache_dict.remove(key); - self.ttl_dict.remove(key); - } -} - -#[cfg(test)] -mod tests { - use std::sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }; - - use super::InMemoryCache; - use std::time::Duration; - - fn cache(now: Arc, max_size: usize, default_ttl: Duration) -> InMemoryCache { - InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || { - Duration::from_secs(now.load(Ordering::Relaxed)) - }) - } - - #[test] - fn ttl_expiry_is_deterministic() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); - cache.set_cache("key", "value".to_string(), None); - assert_eq!(cache.get_cache("key"), Some("value".to_string())); - now.store(161, Ordering::Relaxed); - assert_eq!(cache.get_cache("key"), None); - assert_eq!(cache.get_ttl("key"), None); - } - - #[test] - fn default_and_per_set_ttl_are_applied() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); - cache.set_cache("default", "value".to_string(), None); - cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20))); - assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160))); - assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120))); - } - - #[test] - fn unexpired_entries_do_not_allow_ttl_override() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); - cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20))); - cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80))); - assert_eq!(cache.get_cache("key"), Some("second".to_string())); - assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120))); - now.store(121, Ordering::Relaxed); - cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80))); - assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201))); - } - - #[test] - fn max_size_evicts_earliest_expiration() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now, 2, Duration::from_secs(60)); - cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10))); - cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20))); - cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); - assert_eq!(cache.get_cache("early"), None); - assert!(cache.get_cache("late").is_some()); - assert!(cache.get_cache("new").is_some()); - } - - #[test] - fn expired_entries_are_evicted_before_live_entries() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now.clone(), 3, Duration::from_secs(60)); - cache.set_cache( - "expired-one", - "value".to_string(), - Some(Duration::from_secs(10)), - ); - cache.set_cache( - "expired-two", - "value".to_string(), - Some(Duration::from_secs(20)), - ); - cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100))); - now.store(121, Ordering::Relaxed); - cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100))); - assert_eq!(cache.get_cache("expired-one"), None); - assert_eq!(cache.get_cache("expired-two"), None); - assert!(cache.get_cache("live").is_some()); - assert!(cache.get_cache("new").is_some()); - } - - #[test] - fn stale_heap_entries_are_skipped() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now, 1, Duration::from_secs(60)); - cache.set_cache( - "removed", - "value".to_string(), - Some(Duration::from_secs(10)), - ); - cache.delete_cache("removed"); - cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20))); - cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); - assert_eq!(cache.get_cache("removed"), None); - assert_eq!(cache.get_cache("kept"), None); - assert!(cache.get_cache("new").is_some()); - } - - #[test] - fn delete_and_flush_remove_values_and_ttls() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now, 10, Duration::from_secs(60)); - cache.set_cache("one", "value".to_string(), None); - cache.set_cache("two", "value".to_string(), None); - cache.delete_cache("one"); - assert_eq!(cache.get_cache("one"), None); - cache.flush_cache(); - assert!(cache.cache_dict.is_empty()); - assert!(cache.ttl_dict.is_empty()); - assert!(cache.expiration_heap.is_empty()); - } - - #[test] - fn zero_max_size_does_not_cache() { - let now = Arc::new(AtomicU64::new(100)); - let mut cache = cache(now, 0, Duration::from_secs(60)); - cache.set_cache("key", "value".to_string(), None); - assert_eq!(cache.get_cache("key"), None); - assert!(cache.cache_dict.is_empty()); - } -} diff --git a/litellm-rust/crates/core/src/caching/mod.rs b/litellm-rust/crates/core/src/caching/mod.rs deleted file mode 100644 index 5fb8a0e5174..00000000000 --- a/litellm-rust/crates/core/src/caching/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod in_memory_cache; diff --git a/litellm-rust/crates/core/src/call_lifecycle/cache.rs b/litellm-rust/crates/core/src/call_lifecycle/cache.rs new file mode 100644 index 00000000000..145c4026ebe --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/cache.rs @@ -0,0 +1,71 @@ +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache::{BaseCache, CacheControls, CacheEntry, CacheKwargs}; +use serde_json::Value; + +use super::admission::AdmissionDecline; + +#[derive(Clone, Default)] +pub struct ResponseCachePlan { + pub controls: CacheControls, + pub key: String, + pub backend: Option>>, + pub ttl: Option, + pub max_age: Option, +} + +impl ResponseCachePlan { + pub fn admit(&self) -> Result<(), AdmissionDecline> { + if !self.controls.native_backend && (self.controls.reads() || self.controls.writes()) { + return Err(AdmissionDecline::Feature( + "response cache backend is not implemented in Rust", + )); + } + Ok(()) + } +} + +pub async fn lookup(plan: &ResponseCachePlan) -> Option { + let backend = plan.backend.as_ref().filter(|_| plan.controls.reads())?; + let entry = backend + .async_get_cache(&plan.key, &CacheKwargs::default()) + .await + .ok()??; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + if !entry.fresh(now, plan.max_age) { + return None; + } + match entry.response { + Value::String(text) => serde_json::from_str(&text).ok(), + value => Some(value), + } +} + +pub async fn store(plan: &ResponseCachePlan, value: Option) { + let (Some(value), Some(backend)) = ( + value, + plan.backend.as_ref().filter(|_| plan.controls.writes()), + ) else { + return; + }; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64(); + let _ = backend + .async_set_cache( + &plan.key, + CacheEntry { + timestamp, + response: value, + }, + CacheKwargs { + ttl: plan.ttl, + ..Default::default() + }, + ) + .await; +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/execution.rs b/litellm-rust/crates/core/src/call_lifecycle/execution.rs new file mode 100644 index 00000000000..7a1b9fbeab4 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/execution.rs @@ -0,0 +1,220 @@ +use std::future::Future; + +use tokio::sync::{mpsc, oneshot}; + +use super::host::HostCallStep; +use crate::Error; + +struct PendingOperation { + operation: O, + reply: oneshot::Sender, +} + +pub struct HostExchange { + sender: mpsc::UnboundedSender>, +} + +impl Clone for HostExchange { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + } + } +} + +impl std::fmt::Debug for HostExchange { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("HostExchange") + .finish_non_exhaustive() + } +} + +impl HostExchange { + pub async fn invoke(&self, operation: O) -> Result { + let (reply, receiver) = oneshot::channel(); + self.sender + .send(PendingOperation { operation, reply }) + .map_err(|_| Error::InvalidRequest("host driver was abandoned".into()))?; + receiver + .await + .map_err(|_| Error::InvalidRequest("host operation was abandoned".into())) + } +} + +pub struct HostExecution { + exchange: HostExchange, + operations: mpsc::UnboundedReceiver>, + pending: Option>, + task: Option>>, + completed: bool, + accepts: fn(&O, &R) -> bool, +} + +impl HostExecution { + pub fn new(accepts: fn(&O, &R) -> bool) -> Self { + let (sender, operations) = mpsc::unbounded_channel(); + Self { + exchange: HostExchange { sender }, + operations, + pending: None, + task: None, + completed: false, + accepts, + } + } + + pub fn exchange(&self) -> HostExchange { + self.exchange.clone() + } + + pub fn started(&self) -> bool { + self.task.is_some() || self.completed + } + + pub fn start( + &mut self, + future: impl Future> + Send + 'static, + ) -> Result<(), Error> { + if self.started() { + return Err(Error::InvalidRequest( + "host execution already started".into(), + )); + } + self.task = Some(tokio::spawn(future)); + Ok(()) + } + + pub async fn resume(&mut self, result: Option) -> Result, Error> { + if self.completed { + return Err(Error::InvalidRequest( + "call cannot be resumed after completion".into(), + )); + } + match (&self.pending, &result) { + (Some(pending), Some(reply)) if (self.accepts)(&pending.operation, reply) => {} + (None, None) => {} + _ => { + return Err(Error::InvalidRequest( + "host reply does not match pending operation".into(), + )); + } + } + if let (Some(pending), Some(reply)) = (self.pending.take(), result) { + pending + .reply + .send(reply) + .map_err(|_| Error::InvalidRequest("host operation was abandoned".into()))?; + } + let task = self + .task + .as_mut() + .ok_or_else(|| Error::InvalidRequest("host execution has not started".into()))?; + tokio::select! { + operation = self.operations.recv() => { + let pending = operation.ok_or_else(|| Error::InvalidRequest("host operation channel closed".into()))?; + let operation = pending.operation.clone(); + self.pending = Some(pending); + Ok(HostCallStep::Host(operation)) + } + result = task => { + self.task = None; + self.completed = true; + result.map_err(|error| Error::Network(format!("execution task failed: {error}")))?.map(HostCallStep::Complete) + } + } + } + + pub fn cancel(&mut self) { + self.pending = None; + if let Some(task) = &self.task { + task.abort(); + } + } + + pub async fn stop(&mut self) { + self.cancel(); + if let Some(task) = self.task.as_mut() { + let _ = task.await; + } + self.task = None; + self.completed = true; + } +} + +impl Drop for HostExecution { + fn drop(&mut self) { + if let Some(task) = &self.task { + task.abort(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + + #[tokio::test] + async fn wrong_and_missing_replies_preserve_the_pending_exchange() { + let mut execution = + HostExecution::new(|operation: &u32, reply: &u32| *reply == *operation + 1); + let exchange = execution.exchange(); + execution + .start(async move { exchange.invoke(40).await }) + .unwrap(); + assert!(matches!( + execution.resume(None).await.unwrap(), + HostCallStep::Host(40) + )); + assert!(execution.resume(Some(99)).await.is_err()); + assert!(execution.resume(None).await.is_err()); + assert!(matches!( + execution.resume(Some(41)).await.unwrap(), + HostCallStep::Complete(41) + )); + assert!(execution.resume(None).await.is_err()); + } + + struct Capture(Arc); + impl Drop for Capture { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn stop_waits_for_provider_captures_to_drop() { + let dropped = Arc::new(AtomicBool::new(false)); + let capture = Capture(dropped.clone()); + let mut execution = HostExecution::::new(|_, _| true); + let exchange = execution.exchange(); + execution + .start(async move { + let _capture = capture; + exchange.invoke(1).await + }) + .unwrap(); + assert!(matches!( + execution.resume(None).await.unwrap(), + HostCallStep::Host(1) + )); + execution.stop().await; + assert!(dropped.load(Ordering::SeqCst)); + assert!(execution.resume(Some(2)).await.is_err()); + } + + #[tokio::test] + async fn provider_panics_are_terminal_errors() { + let mut execution = HostExecution::::new(|_, _| true); + execution.start(async { panic!("provider panic") }).unwrap(); + assert!(matches!( + execution.resume(None).await, + Err(Error::Network(_)) + )); + assert!(execution.resume(None).await.is_err()); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/host.rs b/litellm-rust/crates/core/src/call_lifecycle/host.rs index ac6ddf99b9e..98be31d17cb 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/host.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/host.rs @@ -25,6 +25,26 @@ pub trait HostCall: Send + Sync { ) -> HostCallFuture<'_, Self::Operation, Self::Complete>; } +pub type LifecycleBackendFuture<'a, R> = Pin + Send + 'a>>; + +pub trait LifecycleBackend: Send + Sync { + fn invoke(&self, operation: O) -> LifecycleBackendFuture<'_, R>; +} + +pub async fn drive(call: &mut C, backend: &B) -> Result +where + C: HostCall, + B: LifecycleBackend + ?Sized, +{ + let mut result = None; + loop { + match call.resume(result.take()).await? { + HostCallStep::Host(operation) => result = Some(backend.invoke(operation).await), + HostCallStep::Complete(response) => return Ok(response), + } + } +} + pub enum HostStep { Ready(V), Suspend(S), @@ -35,9 +55,12 @@ pub enum HostPhase { Setup, DeploymentPreCall, Prepare, + CacheLookup, Execute, ConstructResponse, + PostProcess, DeploymentPostCall, + CacheStore, Finalize, Success, MapFailure, @@ -56,6 +79,7 @@ pub enum HostFailure { pub struct HostLifecycle { phase: HostPhase, asynchronous: bool, + cached: bool, } impl HostLifecycle { @@ -63,9 +87,15 @@ impl HostLifecycle { Self { phase: HostPhase::Setup, asynchronous, + cached: false, } } + pub fn cache_hit(&mut self) { + self.cached = true; + self.phase = HostPhase::ConstructResponse; + } + pub fn phase(&self) -> HostPhase { self.phase } @@ -104,10 +134,16 @@ impl HostLifecycle { self.phase = match self.phase { HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall, HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare, - HostPhase::Prepare => HostPhase::Execute, + HostPhase::Prepare => HostPhase::CacheLookup, + HostPhase::CacheLookup => HostPhase::Execute, HostPhase::Execute => HostPhase::ConstructResponse, - HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall, - HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize, + HostPhase::ConstructResponse => HostPhase::PostProcess, + HostPhase::PostProcess if self.cached => HostPhase::Finalize, + HostPhase::PostProcess if self.asynchronous => HostPhase::DeploymentPostCall, + HostPhase::PostProcess => HostPhase::CacheStore, + HostPhase::DeploymentPostCall if self.cached => HostPhase::Finalize, + HostPhase::DeploymentPostCall => HostPhase::CacheStore, + HostPhase::CacheStore => HostPhase::Finalize, HostPhase::Finalize => HostPhase::Success, HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure, HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure, diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index e85167ab14c..86b28def804 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -1,420 +1,13 @@ -use std::future::Future; -use std::time::{Instant, SystemTime, UNIX_EPOCH}; - -use crate::Error; - pub mod admission; +pub mod cache; pub mod dispatch; +pub mod execution; pub mod host; #[cfg(test)] #[path = "../../tests/host_lifecycle.rs"] mod host_tests; +pub mod provider; pub mod types; +pub mod workflow; -pub use types::{ - CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest, - CallLifecycleTiming, -}; - -pub trait CallLifecycleHooks: Send + Sync { - type PreCallFuture<'a>: Future> + Send + 'a - where - Self: 'a, - InitialReq: 'a, - ProviderReq: 'a, - Resp: 'a; - - type DuringCallFuture<'a>: Future> + Send + 'a - where - Self: 'a, - InitialReq: 'a, - ProviderReq: 'a, - Resp: 'a; - - type SuccessFuture<'a>: Future + Send + 'a - where - Self: 'a, - Resp: 'a; - - type FailureFuture<'a>: Future + Send + 'a - where - Self: '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>; - - fn async_log_success_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - response: &'a Resp, - timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a>; - - fn async_log_failure_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - error: &'a Error, - timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a>; -} - -pub trait CallLifecycleObserver: Send + Sync { - fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {} - - fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {} -} - -#[derive(Default)] -pub struct NoopCallLifecycleObserver; - -impl CallLifecycleObserver for NoopCallLifecycleObserver {} - -pub struct CallLifecycle<'a> { - observer: &'a dyn CallLifecycleObserver, -} - -impl<'a> CallLifecycle<'a> { - pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self { - Self { observer } - } - - pub async fn run_request( - &self, - request: InitialReq, - hooks: &Hooks, - provider_call: ProviderCall, - ) -> Result - where - InitialReq: CallLifecycleRequest, - Hooks: CallLifecycleHooks, - ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, - { - let context = request.lifecycle_context(); - self.run(context, request, hooks, provider_call).await - } - - pub async fn run( - &self, - context: CallLifecycleContext, - request: InitialReq, - hooks: &Hooks, - provider_call: ProviderCall, - ) -> Result - where - Hooks: CallLifecycleHooks, - ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, - { - let call_start = epoch_seconds(); - let mut phases = Vec::new(); - - let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall); - let request = match hooks.async_pre_call_hook(&context, request).await { - Ok(request) => { - phases.push(self.finish_phase(&context, pre_call)); - request - } - Err(error) => { - phases.push(self.finish_phase(&context, pre_call)); - self.log_failure(&context, hooks, &error, call_start, &mut phases) - .await; - return Err(error); - } - }; - - let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall); - let provider_request = match hooks.async_during_call_hook(&context, request).await { - Ok(request) => { - phases.push(self.finish_phase(&context, during_call)); - request - } - Err(error) => { - phases.push(self.finish_phase(&context, during_call)); - self.log_failure(&context, hooks, &error, call_start, &mut phases) - .await; - return Err(error); - } - }; - - let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall); - let result = provider_call(provider_request).await; - phases.push(self.finish_phase(&context, provider_phase)); - - match &result { - Ok(response) => { - let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback); - let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); - hooks - .async_log_success_event(&context, response, &timing) - .await; - phases.push(self.finish_phase(&context, success_phase)); - } - Err(error) => { - self.log_failure(&context, hooks, error, call_start, &mut phases) - .await; - } - } - - result - } - - async fn log_failure( - &self, - context: &CallLifecycleContext, - hooks: &Hooks, - error: &Error, - call_start: f64, - phases: &mut Vec, - ) where - Hooks: CallLifecycleHooks, - { - let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback); - let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); - hooks.async_log_failure_event(context, error, &timing).await; - phases.push(self.finish_phase(context, failure_phase)); - } - - fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart { - self.observer.on_phase_start(context, phase); - PhaseStart { - phase, - start_time: epoch_seconds(), - started_at: Instant::now(), - } - } - - fn finish_phase( - &self, - context: &CallLifecycleContext, - phase_start: PhaseStart, - ) -> CallLifecyclePhaseTiming { - let timing = CallLifecyclePhaseTiming { - phase: phase_start.phase, - start_time: phase_start.start_time, - end_time: epoch_seconds(), - duration: phase_start.started_at.elapsed(), - }; - self.observer.on_phase_end(context, &timing); - timing - } -} - -impl Default for CallLifecycle<'static> { - fn default() -> Self { - static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver; - Self::new(&OBSERVER) - } -} - -struct PhaseStart { - phase: CallLifecyclePhase, - start_time: f64, - started_at: Instant, -} - -fn epoch_seconds() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs_f64()) - .unwrap_or(0.0) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::pin::Pin; - use std::sync::Mutex; - - type BoxFuture<'a, T> = Pin + Send + 'a>>; - - #[derive(Default)] - struct RecordingHooks { - events: Mutex>, - } - - struct RecordingRequest(String); - - impl CallLifecycleRequest for RecordingRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1") - } - } - - impl RecordingHooks { - fn events(&self) -> Vec<&'static str> { - self.events.lock().unwrap().clone() - } - } - - impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, Result>; - type DuringCallFuture<'a> = BoxFuture<'a, Result>; - type SuccessFuture<'a> = BoxFuture<'a, ()>; - type FailureFuture<'a> = BoxFuture<'a, ()>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: String, - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("pre_call"); - Ok(format!("{request}:pre")) - }) - } - - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: String, - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("during_call"); - Ok(format!("{request}:during")) - }) - } - - fn async_log_success_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _response: &'a String, - timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - Box::pin(async move { - assert!(timing.end_time >= timing.start_time); - assert_eq!(timing.phases.len(), 3); - self.events.lock().unwrap().push("success"); - }) - } - - fn async_log_failure_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _error: &'a Error, - _timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("failure"); - }) - } - } - - impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, Result>; - type DuringCallFuture<'a> = BoxFuture<'a, Result>; - type SuccessFuture<'a> = BoxFuture<'a, ()>; - type FailureFuture<'a> = BoxFuture<'a, ()>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: RecordingRequest, - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("pre_call"); - Ok(RecordingRequest(format!("{}:pre", request.0))) - }) - } - - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: RecordingRequest, - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("during_call"); - Ok(format!("{}:during", request.0)) - }) - } - - fn async_log_success_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _response: &'a String, - _timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("success"); - }) - } - - fn async_log_failure_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _error: &'a Error, - _timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("failure"); - }) - } - } - - #[tokio::test] - async fn lifecycle_runs_hooks_around_provider_call() { - let hooks = RecordingHooks::default(); - let response = CallLifecycle::default() - .run( - CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), - "request".to_string(), - &hooks, - |request| async move { - assert_eq!(request, "request:pre:during"); - Ok("response".to_string()) - }, - ) - .await - .expect("call succeeds"); - - assert_eq!(response, "response"); - assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); - } - - #[tokio::test] - async fn lifecycle_logs_failure_when_provider_fails() { - let hooks = RecordingHooks::default(); - let error = CallLifecycle::default() - .run( - CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), - "request".to_string(), - &hooks, - |_request| async move { - Err::(Error::Network("provider down".to_string())) - }, - ) - .await - .expect_err("call fails"); - - assert_eq!(error, Error::Network("provider down".to_string())); - assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); - } - - #[tokio::test] - async fn lifecycle_can_run_any_request_with_embedded_context() { - let hooks = RecordingHooks::default(); - let response = CallLifecycle::default() - .run_request( - RecordingRequest("request".to_string()), - &hooks, - |request| async move { - assert_eq!(request, "request:pre:during"); - Ok("response".to_string()) - }, - ) - .await - .expect("call succeeds"); - - assert_eq!(response, "response"); - assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); - } -} +pub use types::{CallLifecycleContext, CallLifecycleRequest, CallLifecycleTiming}; diff --git a/litellm-rust/crates/core/src/call_lifecycle/provider.rs b/litellm-rust/crates/core/src/call_lifecycle/provider.rs new file mode 100644 index 00000000000..f1d96d473b6 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/provider.rs @@ -0,0 +1,516 @@ +use std::future::Future; +use std::marker::PhantomData; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use super::cache::{ResponseCachePlan, lookup, store}; +use super::execution::HostExchange; +use super::host::{HostFailure, LifecycleBackend, LifecycleBackendFuture, drive}; +use super::workflow::{LifecycleCall, LifecycleOperation, Workflow, WorkflowFuture, WorkflowReply}; +use crate::Error; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; + +#[derive(Clone, Default)] +pub struct ProviderOptions { + pub model: String, + pub litellm_call_id: Option, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub timeout: Option, +} + +impl ProviderOptions { + pub fn lifecycle_context(&self, call_type: &str) -> CallLifecycleContext { + CallLifecycleContext::new( + call_type, + self.model.clone(), + self.custom_llm_provider.clone().unwrap_or_default(), + self.litellm_call_id + .clone() + .unwrap_or_else(|| format!("native-{:032x}", rand::random::())), + ) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ProviderRequest { + pub model: String, + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Value, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ProviderResponse { + pub status: u16, + pub body: String, +} + +pub type ProviderHookFuture<'a, T> = Pin> + Send + 'a>>; + +pub trait ProviderHooks: Send + Sync { + fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest>; + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse>; +} + +pub struct NoopProviderHooks; + +impl ProviderHooks for NoopProviderHooks { + fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> { + Box::pin(async move { Ok(request) }) + } + + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse> { + Box::pin(async move { Ok(response) }) + } +} + +pub struct ProviderHookChain { + hooks: Box<[Arc]>, +} + +impl ProviderHookChain { + pub fn new(hooks: impl IntoIterator>) -> Self { + Self { + hooks: hooks.into_iter().collect(), + } + } + + fn before<'a>( + hooks: &'a [Arc], + request: ProviderRequest, + ) -> ProviderHookFuture<'a, ProviderRequest> { + Box::pin(async move { + let Some((hook, remaining)) = hooks.split_first() else { + return Ok(request); + }; + let request = hook.before_request(request).await?; + Self::before(remaining, request).await + }) + } + + fn after<'a>( + hooks: &'a [Arc], + response: ProviderResponse, + ) -> ProviderHookFuture<'a, ProviderResponse> { + Box::pin(async move { + let Some((hook, remaining)) = hooks.split_first() else { + return Ok(response); + }; + let response = hook.after_response(response).await?; + Self::after(remaining, response).await + }) + } +} + +impl ProviderHooks for ProviderHookChain { + fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> { + Self::before(&self.hooks, request) + } + + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse> { + Self::after(&self.hooks, response) + } +} + +#[derive(Clone, Debug)] +pub enum CompletedOperation { + Lifecycle(LifecycleOperation), + BeforeRequest(ProviderRequest), + AfterResponse(ProviderResponse), +} + +pub enum CompletedReply { + Prepared(Result), + CacheStore(Result, Error>), + Request(Result), + Lifecycle(Result<(), HostFailure>), + BeforeRequest(Result), + AfterResponse(Result), +} + +pub trait CompletedRoute: Send + Sync + 'static { + type Request: Send + Sync + 'static; + type Response: Clone + Send + Sync + serde::de::DeserializeOwned + 'static; + + fn run(request: Self::Request, hooks: Arc) + -> WorkflowFuture; + fn context(request: &Self::Request) -> CallLifecycleContext; +} + +pub struct CompletedWorkflow { + route: PhantomData, + hooks: Arc, + cache: ResponseCachePlan, + cached: Arc>>, + pending_write: Option, + terminal: Arc>>, +} + +impl Default for CompletedWorkflow { + fn default() -> Self { + Self { + route: PhantomData, + hooks: Arc::new(NoopProviderHooks), + cache: ResponseCachePlan::default(), + cached: Arc::default(), + pending_write: None, + terminal: Arc::default(), + } + } +} + +impl CompletedWorkflow { + pub fn with_hooks(hooks: Arc) -> Self { + Self { + hooks, + ..Self::default() + } + } +} + +impl Workflow for CompletedWorkflow { + type Request = R::Request; + type Operation = CompletedOperation; + type Reply = CompletedReply; + type Response = R::Response; + + fn operation(operation: LifecycleOperation) -> Self::Operation { + CompletedOperation::Lifecycle(operation) + } + + fn accepts(operation: &Self::Operation, reply: &Self::Reply) -> bool { + match (operation, reply) { + (_, CompletedReply::Lifecycle(Err(_))) => true, + ( + CompletedOperation::Lifecycle(LifecycleOperation::Phase( + super::host::HostPhase::Prepare, + )), + CompletedReply::Prepared(_), + ) => true, + ( + CompletedOperation::Lifecycle(LifecycleOperation::Phase( + super::host::HostPhase::CacheStore, + )), + CompletedReply::CacheStore(_), + ) => true, + ( + CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest), + CompletedReply::Request(_), + ) => true, + (CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest), _) => false, + (CompletedOperation::Lifecycle(_), CompletedReply::Lifecycle(_)) => true, + (CompletedOperation::BeforeRequest(_), CompletedReply::BeforeRequest(_)) => true, + (CompletedOperation::AfterResponse(_), CompletedReply::AfterResponse(_)) => true, + _ => false, + } + } + + fn reply(&mut self, reply: Self::Reply) -> WorkflowReply { + match reply { + CompletedReply::Prepared(result) => WorkflowReply::Lifecycle( + result + .map(|cache| self.cache = cache) + .map_err(HostFailure::Error), + ), + CompletedReply::CacheStore(result) => WorkflowReply::Lifecycle( + result + .map(|value| self.pending_write = value) + .map_err(HostFailure::Error), + ), + CompletedReply::Request(request) => WorkflowReply::Request(request), + CompletedReply::Lifecycle(result) => WorkflowReply::Lifecycle(result), + reply => WorkflowReply::Operation(reply), + } + } + + fn start( + &mut self, + request: Self::Request, + host: HostExchange, + ) -> WorkflowFuture { + let context = R::context(&request); + let started = epoch_seconds(); + let terminal = self.terminal.clone(); + let host_hooks: Arc = Arc::new(ExchangeHooks { host }); + let future = R::run( + request, + Arc::new(ProviderHookChain::new([self.hooks.clone(), host_hooks])), + ); + Box::pin(async move { + let result = future.await; + let timing = CallLifecycleTiming::new(started, epoch_seconds()); + *terminal.lock().unwrap_or_else(|error| error.into_inner()) = Some((context, timing)); + result + }) + } + + fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> { + self.terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + + fn cache_lookup(&mut self) -> WorkflowFuture> { + let cached = self.cached.clone(); + let cache = self.cache.clone(); + Box::pin(async move { + let Some(value) = lookup(&cache).await else { + return Ok(None); + }; + let Ok(response) = serde_json::from_value(value.clone()) else { + return Ok(None); + }; + *cached.lock().unwrap_or_else(|error| error.into_inner()) = Some(value); + Ok(Some(response)) + }) + } + + fn cached_public_response(&self) -> Option { + self.cached + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + + fn flush_cache(&mut self) -> WorkflowFuture<()> { + let value = self.pending_write.take(); + let cache = self.cache.clone(); + Box::pin(async move { + store(&cache, value).await; + Ok(()) + }) + } +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +struct ExchangeHooks { + host: HostExchange, CompletedReply>, +} + +impl ProviderHooks for ExchangeHooks { + fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> { + Box::pin(async move { + match self + .host + .invoke(CompletedOperation::BeforeRequest(request)) + .await? + { + CompletedReply::BeforeRequest(result) => result, + _ => Err(Error::InvalidRequest( + "unexpected provider request reply".into(), + )), + } + }) + } + + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse> { + Box::pin(async move { + match self + .host + .invoke(CompletedOperation::AfterResponse(response)) + .await? + { + CompletedReply::AfterResponse(result) => result, + _ => Err(Error::InvalidRequest( + "unexpected provider response reply".into(), + )), + } + }) + } +} + +pub type CompletedCall = LifecycleCall>; + +pub async fn run_completed(request: R::Request) -> Result { + run_completed_with_hooks::(request, Arc::new(NoopProviderHooks)).await +} + +pub async fn run_completed_with_hooks( + request: R::Request, + hooks: Arc, +) -> Result { + let mut call = CompletedCall::::new(CompletedWorkflow::with_hooks(hooks), false); + drive( + &mut call, + &CompletedBackend:: { + request: Mutex::new(Some(request)), + route: PhantomData, + }, + ) + .await +} + +struct CompletedBackend { + request: Mutex>, + route: PhantomData, +} + +impl + LifecycleBackend, CompletedReply> + for CompletedBackend +{ + fn invoke( + &self, + operation: CompletedOperation, + ) -> LifecycleBackendFuture<'_, CompletedReply> { + Box::pin(async move { + match operation { + CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => { + CompletedReply::Request( + self.request + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .ok_or_else(|| { + Error::InvalidRequest("request was already projected".into()) + }), + ) + } + CompletedOperation::Lifecycle(_) => CompletedReply::Lifecycle(Ok(())), + CompletedOperation::BeforeRequest(request) => { + CompletedReply::BeforeRequest(Ok(request)) + } + CompletedOperation::AfterResponse(response) => { + CompletedReply::AfterResponse(Ok(response)) + } + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct RecordingHook { + name: &'static str, + events: Arc>>, + } + + impl ProviderHooks for RecordingHook { + fn before_request( + &self, + request: ProviderRequest, + ) -> ProviderHookFuture<'_, ProviderRequest> { + Box::pin(async move { + self.events + .lock() + .unwrap() + .push(format!("{}:before", self.name)); + Ok(ProviderRequest { + body: Value::String(format!( + "{}:{}", + request.body.as_str().unwrap(), + self.name + )), + ..request + }) + }) + } + + fn after_response( + &self, + response: ProviderResponse, + ) -> ProviderHookFuture<'_, ProviderResponse> { + Box::pin(async move { + self.events + .lock() + .unwrap() + .push(format!("{}:after", self.name)); + Ok(ProviderResponse { + body: format!("{}:{}", response.body, self.name), + ..response + }) + }) + } + } + + #[derive(Clone, Deserialize)] + struct TestResponse(String); + + struct TestRoute; + + impl CompletedRoute for TestRoute { + type Request = (); + type Response = TestResponse; + + fn run((): Self::Request, hooks: Arc) -> WorkflowFuture { + Box::pin(async move { + let request = hooks + .before_request(ProviderRequest { + model: "model".into(), + url: "https://example.com".into(), + headers: Vec::new(), + body: Value::String("request".into()), + }) + .await?; + let response = hooks + .after_response(ProviderResponse { + status: 200, + body: request.body.as_str().unwrap().to_owned(), + }) + .await?; + Ok(TestResponse(response.body)) + }) + } + + fn context(_: &Self::Request) -> CallLifecycleContext { + CallLifecycleContext::new("test", "model", "provider", "call") + } + } + + #[tokio::test] + async fn native_hook_chain_runs_in_order_through_the_lifecycle() { + let events = Arc::new(Mutex::new(Vec::new())); + let hooks = ProviderHookChain::new([ + Arc::new(RecordingHook { + name: "first", + events: events.clone(), + }) as Arc, + Arc::new(RecordingHook { + name: "second", + events: events.clone(), + }) as Arc, + ]); + let response = run_completed_with_hooks::((), Arc::new(hooks)) + .await + .unwrap(); + assert_eq!(response.0, "request:first:second:first:second"); + assert_eq!( + *events.lock().unwrap(), + [ + "first:before", + "second:before", + "first:after", + "second:after" + ] + ); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/types.rs b/litellm-rust/crates/core/src/call_lifecycle/types.rs index 8819c8830d2..891133792ae 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/types.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/types.rs @@ -1,5 +1,3 @@ -use std::time::Duration; - #[derive(Clone, Debug, PartialEq, Eq)] pub struct CallLifecycleContext { pub call_type: String, @@ -28,48 +26,17 @@ pub trait CallLifecycleRequest { fn lifecycle_context(&self) -> CallLifecycleContext; } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum CallLifecyclePhase { - PreCall, - DuringCall, - ProviderCall, - SuccessCallback, - FailureCallback, -} - -impl CallLifecyclePhase { - pub fn as_str(self) -> &'static str { - match self { - Self::PreCall => "pre_call", - Self::DuringCall => "during_call", - Self::ProviderCall => "provider_call", - Self::SuccessCallback => "success_callback", - Self::FailureCallback => "failure_callback", - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct CallLifecyclePhaseTiming { - pub phase: CallLifecyclePhase, - pub start_time: f64, - pub end_time: f64, - pub duration: Duration, -} - #[derive(Clone, Debug, PartialEq)] pub struct CallLifecycleTiming { pub start_time: f64, pub end_time: f64, - pub phases: Vec, } impl CallLifecycleTiming { - pub fn new(start_time: f64, end_time: f64, phases: Vec) -> Self { + pub fn new(start_time: f64, end_time: f64) -> Self { Self { start_time, end_time, - phases, } } } diff --git a/litellm-rust/crates/core/src/call_lifecycle/workflow.rs b/litellm-rust/crates/core/src/call_lifecycle/workflow.rs new file mode 100644 index 00000000000..c1a87a5d41c --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/workflow.rs @@ -0,0 +1,250 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use super::execution::{HostExchange, HostExecution}; +use super::host::{HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase}; +use super::{CallLifecycleContext, CallLifecycleTiming}; +use crate::Error; + +pub type WorkflowFuture = Pin> + Send + 'static>>; + +#[derive(Clone, Debug)] +pub enum LifecycleOperation { + ProjectRequest, + Phase(HostPhase), + ConstructResponse(Arc), + ConstructCachedResponse(serde_json::Value), + MapFailure(Error), + Success { + context: CallLifecycleContext, + response: Arc, + timing: CallLifecycleTiming, + }, + Failure { + context: CallLifecycleContext, + error: Error, + timing: CallLifecycleTiming, + }, +} + +pub enum WorkflowReply { + Request(Result), + Lifecycle(Result<(), HostFailure>), + Operation(R), +} + +pub trait Workflow: Send + Sync { + type Request: Send + 'static; + type Operation: Clone + Send + Sync + 'static; + type Reply: Send + Sync + 'static; + type Response: Clone + Send + Sync + 'static; + + fn operation(operation: LifecycleOperation) -> Self::Operation; + fn accepts(operation: &Self::Operation, reply: &Self::Reply) -> bool; + fn reply(&mut self, reply: Self::Reply) -> WorkflowReply; + fn start( + &mut self, + request: Self::Request, + host: HostExchange, + ) -> WorkflowFuture; + fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> { + None + } + fn cache_lookup(&mut self) -> WorkflowFuture> { + Box::pin(async { Ok(None) }) + } + fn cached_public_response(&self) -> Option { + None + } + fn flush_cache(&mut self) -> WorkflowFuture<()> { + Box::pin(async { Ok(()) }) + } +} + +pub struct LifecycleCall { + workflow: W, + lifecycle: HostLifecycle, + execution: HostExecution, + pending: Option, + response: Option>, + error: Option, + completed: bool, +} + +impl LifecycleCall { + pub fn new(workflow: W, asynchronous: bool) -> Self { + Self { + workflow, + lifecycle: HostLifecycle::new(asynchronous), + execution: HostExecution::new(W::accepts), + pending: None, + response: None, + error: None, + completed: false, + } + } + + pub async fn resume( + &mut self, + result: Option, + ) -> Result, Error> { + if self.completed { + return Err(Error::InvalidRequest( + "call cannot be resumed after completion".into(), + )); + } + match (&self.pending, &result) { + (Some(operation), Some(reply)) if W::accepts(operation, reply) => {} + (None, None) => {} + _ => { + return Err(Error::InvalidRequest( + "host reply does not match pending operation".into(), + )); + } + } + self.pending = None; + let provider_reply = match result.map(|reply| self.workflow.reply(reply)) { + Some(WorkflowReply::Request(Ok(request))) => { + let future = self.workflow.start(request, self.execution.exchange()); + self.execution.start(future)?; + None + } + Some(WorkflowReply::Request(Err(error))) => { + self.accept(Err(HostFailure::Error(error))); + None + } + Some(WorkflowReply::Lifecycle(result)) => { + self.accept(result); + None + } + Some(WorkflowReply::Operation(reply)) => Some(reply), + None => None, + }; + self.workflow.flush_cache().await?; + if self.lifecycle.phase() == HostPhase::CacheLookup { + match self.workflow.cache_lookup().await { + Ok(Some(response)) => { + self.response = Some(Arc::new(response)); + self.lifecycle.cache_hit(); + } + Ok(None) => self.accept(Ok(())), + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + } + if self.lifecycle.phase() == HostPhase::Execute { + if !self.execution.started() { + return Ok(self.host_step(W::operation(LifecycleOperation::ProjectRequest))); + } + match self.execution.resume(provider_reply).await { + Ok(HostCallStep::Host(operation)) => return Ok(self.host_step(operation)), + Ok(HostCallStep::Complete(response)) => { + self.response = Some(Arc::new(response)); + self.accept(Ok(())); + } + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + } + if self.error.is_some() { + self.execution.stop().await; + } + let operation = match self.lifecycle.phase() { + HostPhase::Complete => { + self.completed = true; + return match self.error.take() { + Some(error) => Err(error), + None => self + .response + .take() + .map(Arc::unwrap_or_clone) + .map(HostCallStep::Complete) + .ok_or_else(|| { + Error::InvalidRequest("call completed without a response".into()) + }), + }; + } + HostPhase::ConstructResponse => match self.workflow.cached_public_response() { + Some(response) => LifecycleOperation::ConstructCachedResponse(response), + None => LifecycleOperation::ConstructResponse(self.response()?), + }, + HostPhase::MapFailure => LifecycleOperation::MapFailure(self.error()?), + phase => match (phase, self.workflow.terminal()) { + (HostPhase::Success, Some((context, timing))) => LifecycleOperation::Success { + context, + response: self.response()?, + timing, + }, + (HostPhase::Failure, Some((context, timing))) => LifecycleOperation::Failure { + context, + error: self.error()?, + timing, + }, + _ => LifecycleOperation::Phase(phase), + }, + }; + Ok(self.host_step(W::operation(operation))) + } + + fn response(&self) -> Result, Error> { + self.response + .clone() + .ok_or_else(|| Error::InvalidRequest("missing response".into())) + } + + fn error(&self) -> Result { + self.error + .clone() + .ok_or_else(|| Error::InvalidRequest("missing failure".into())) + } + + fn host_step(&mut self, operation: W::Operation) -> HostCallStep { + self.pending = Some(operation.clone()); + HostCallStep::Host(operation) + } + + fn accept(&mut self, result: Result<(), HostFailure>) { + let cancelled = matches!(&result, Err(HostFailure::Cancelled(_))); + if let Some(error) = self.lifecycle.accept(result) { + if cancelled { + self.error = Some(error); + } else { + self.error.get_or_insert(error); + } + self.execution.cancel(); + } + } + + pub async fn interrupt( + &mut self, + failure: HostFailure, + ) -> Result, Error> { + if self.completed { + return Err(Error::InvalidRequest( + "call cannot be interrupted after completion".into(), + )); + } + self.pending = None; + self.accept(Err(failure)); + self.resume(None).await + } +} + +impl HostCall for LifecycleCall { + type Operation = W::Operation; + type Result = W::Reply; + type Complete = W::Response; + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(LifecycleCall::resume(self, result)) + } + + fn interrupt( + &mut self, + failure: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(LifecycleCall::interrupt(self, failure)) + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 96d001e2892..d1c878cda8e 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,5 +1,6 @@ use serde_json::Value; +use crate::call_lifecycle::provider::{ProviderHooks, ProviderRequest, ProviderResponse}; use crate::error::Error; use crate::http_utils::{http_request, truncate_error_body}; @@ -14,8 +15,24 @@ use super::types::{ #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_chat_completions_provider_call( request: ResolvedChatCompletionsRequest<'_>, + hooks: &dyn ProviderHooks, ) -> Result { let request = prepare_provider_request(request)?; + let changed = hooks + .before_request(ProviderRequest { + model: request.model.clone(), + url: request.url.clone(), + headers: request.upstream_headers.clone(), + body: request.body.clone(), + }) + .await?; + let request = ProviderChatCompletionsRequest { + model: changed.model, + url: changed.url, + upstream_headers: changed.headers, + body: changed.body, + ..request + }; let body = serde_json::to_vec(&request.body).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize chat completions request: {err}" @@ -48,9 +65,17 @@ pub(super) async fn execute_chat_completions_provider_call( .await .map_err(|err| Error::Network(err.to_string()))?; - if !status.is_success() { - return Err(Error::Http { + let observed = hooks + .after_response(ProviderResponse { status: status.as_u16(), + body: text, + }) + .await?; + let observed_status = observed.status; + let text = observed.body; + if !(200..300).contains(&observed_status) { + return Err(Error::Http { + status: observed_status, body: truncate_error_body(&text), }); } diff --git a/litellm-rust/crates/core/src/chat_completions/lifecycle.rs b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs new file mode 100644 index 00000000000..ca6bd37fe47 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use serde_json::Value; + +use super::types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use crate::call_lifecycle::provider::{ + CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions, +}; +use crate::call_lifecycle::workflow::WorkflowFuture; + +pub struct OwnedChatCompletionsRequest { + pub options: ProviderOptions, + pub messages: Value, + pub optional_params: serde_json::Map, +} + +impl From> for OwnedChatCompletionsRequest { + fn from(request: ChatCompletionsRequest<'_>) -> Self { + Self { + options: ProviderOptions { + model: request.model.to_owned(), + litellm_call_id: None, + api_key: request.api_key.map(str::to_owned), + api_base: request.api_base.map(str::to_owned), + custom_llm_provider: request.custom_llm_provider.map(str::to_owned), + extra_headers: request.extra_headers, + timeout: request.timeout, + }, + messages: request.messages, + optional_params: request.optional_params, + } + } +} + +pub struct ChatCompletionsRoute; +pub type ChatCompletionsCall = CompletedCall; + +impl CompletedRoute for ChatCompletionsRoute { + type Request = OwnedChatCompletionsRequest; + type Response = ChatCompletionsResponse; + + fn run( + request: Self::Request, + hooks: Arc, + ) -> WorkflowFuture { + Box::pin(async move { + let options = request.options; + let request = ChatCompletionsRequest { + model: &options.model, + api_key: options.api_key.as_deref(), + api_base: options.api_base.as_deref(), + custom_llm_provider: options.custom_llm_provider.as_deref(), + extra_headers: options.extra_headers, + timeout: options.timeout, + messages: request.messages, + optional_params: request.optional_params, + }; + super::handler::execute_chat_completions_provider_call( + super::prepare::resolve_request(request)?, + hooks.as_ref(), + ) + .await + }) + } + + fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext { + request.options.lifecycle_context("completion") + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 6ac42928dd2..c0f325712c3 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -11,6 +11,7 @@ mod client; mod common_utils; pub mod conversation; pub(crate) mod handler; +pub mod lifecycle; mod prepare; pub mod response_utils; pub mod transformation; @@ -18,15 +19,17 @@ pub mod types; use serde_json::{Map, Value}; -use handler::execute_chat_completions_provider_call; -use prepare::{parse_messages, resolve_provider_config, resolve_request}; +use prepare::{parse_messages, resolve_provider_config}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn chat_completions( request: ChatCompletionsRequest<'_>, ) -> Result { - execute_chat_completions_provider_call(resolve_request(request)?).await + crate::call_lifecycle::provider::run_completed::( + request.into(), + ) + .await } /// Whether the core would accept this request, without resolving credentials or diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 0b3573deab2..921ef4d9b86 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,6 +1,5 @@ pub mod audio_transcription; pub mod auth; -pub mod caching; pub mod call_lifecycle; pub mod chat_completions; pub mod constants; diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 61ff81bcdc8..31c668efaa7 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,3 +1,4 @@ +use crate::call_lifecycle::provider::{ProviderHooks, ProviderRequest, ProviderResponse}; use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; use crate::error::Error; use crate::http_utils::http_request; @@ -10,8 +11,24 @@ use super::types::{AnthropicMessagesResponse, MessagesRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_messages_provider_call( request: MessagesRequest<'_>, + hooks: &dyn ProviderHooks, ) -> Result { let request = prepare_provider_request(request)?; + let changed = hooks + .before_request(ProviderRequest { + model: request.model.clone(), + url: request.url.clone(), + headers: request.upstream_headers.clone(), + body: request.body.clone(), + }) + .await?; + let request = super::types::ProviderMessagesRequest { + model: changed.model, + url: changed.url, + upstream_headers: changed.headers, + body: changed.body, + ..request + }; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -30,9 +47,17 @@ pub(super) async fn execute_messages_provider_call( .await .map_err(|err| Error::Network(err.to_string()))?; - if !status.is_success() { - return Err(Error::Http { + let observed = hooks + .after_response(ProviderResponse { status: status.as_u16(), + body: text, + }) + .await?; + let observed_status = observed.status; + let text = observed.body; + if !(200..300).contains(&observed_status) { + return Err(Error::Http { + status: observed_status, body: truncate_error_body(&text), }); } diff --git a/litellm-rust/crates/core/src/messages/lifecycle.rs b/litellm-rust/crates/core/src/messages/lifecycle.rs new file mode 100644 index 00000000000..07dedb79cf0 --- /dev/null +++ b/litellm-rust/crates/core/src/messages/lifecycle.rs @@ -0,0 +1,62 @@ +use std::sync::Arc; + +use serde_json::Value; + +use super::types::{AnthropicMessagesResponse, MessagesRequest}; +use crate::call_lifecycle::provider::{ + CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions, +}; +use crate::call_lifecycle::workflow::WorkflowFuture; + +pub struct OwnedMessagesRequest { + pub options: ProviderOptions, + pub body: Value, +} + +impl From> for OwnedMessagesRequest { + fn from(request: MessagesRequest<'_>) -> Self { + Self { + options: ProviderOptions { + model: request.model.to_owned(), + litellm_call_id: None, + api_key: request.api_key.map(str::to_owned), + api_base: request.api_base.map(str::to_owned), + custom_llm_provider: request.custom_llm_provider.map(str::to_owned), + extra_headers: request.extra_headers, + timeout: request.timeout, + }, + body: request.body, + } + } +} + +pub struct MessagesRoute; +pub type MessagesCall = CompletedCall; + +impl CompletedRoute for MessagesRoute { + type Request = OwnedMessagesRequest; + type Response = AnthropicMessagesResponse; + + fn run( + request: Self::Request, + hooks: Arc, + ) -> WorkflowFuture { + Box::pin(async move { + let options = request.options; + let request = MessagesRequest { + model: &options.model, + api_key: options.api_key.as_deref(), + api_base: options.api_base.as_deref(), + custom_llm_provider: options.custom_llm_provider.as_deref(), + extra_headers: options.extra_headers, + timeout: options.timeout, + body: request.body, + }; + super::handler::execute_messages_provider_call(request, hooks.as_ref()).await + }) + } + + fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext { + request.options.lifecycle_context("anthropic_messages") + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 02106d4d228..9b817b4b69a 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -11,16 +11,17 @@ use crate::Error; mod client; mod common_utils; mod handler; +pub mod lifecycle; mod prepare; pub mod transformation; pub mod types; -use handler::{execute_messages_provider_call, execute_messages_provider_stream}; +use handler::execute_messages_provider_stream; use types::{AnthropicMessagesResponse, MessagesRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(request).await + crate::call_lifecycle::provider::run_completed::(request.into()).await } pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index cd1d538aaa8..32757d3f799 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,10 +1,9 @@ use super::OcrClient; use super::adapters::OcrAdapter; -use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest}; +use super::hooks::{OcrHooks, OcrPostCallRequest, OcrPreCallRequest}; use super::registry::OcrAdapterKind; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use crate::Error; -use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; use std::sync::Arc; pub(crate) async fn perform_ocr_request( @@ -12,28 +11,38 @@ pub(crate) async fn perform_ocr_request( request: LiteLLMOcrRequest, ) -> Result { request.response_format()?; - let context = CallLifecycleContext::new( - "ocr", - request.model.clone(), - request.adapter.provider().as_str(), - request - .litellm_call_id - .clone() - .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::())), - ); - let hooks = OcrLifecycleHooks { - hooks: request.hooks.clone(), - provider_name: context.custom_llm_provider.clone(), - }; - CallLifecycle::default() - .run(context, request, &hooks, |request| async move { - PreparedOcrCall::prepare(client.clone(), request) - .await? - .execute() - .await? - .normalize() + let request = pre_call(request).await?; + PreparedOcrCall::prepare(client.clone(), request) + .await? + .execute() + .await? + .normalize() +} + +async fn pre_call(request: LiteLLMOcrRequest) -> Result { + if !request.hooks.intercepts_requests() { + return Ok(request); + } + let changed = request + .hooks + .pre_call(OcrPreCallRequest { + model: request.model.clone(), + custom_llm_provider: request.adapter.provider().as_str().to_owned(), + document: request.document, + optional_params: serde_json::Value::Object(request.optional_params), }) - .await + .await?; + let serde_json::Value::Object(optional_params) = changed.optional_params else { + return Err(super::error::OcrRequestError::RequestField { + path: "guardrail.optional_params".into(), + } + .into()); + }; + Ok(LiteLLMOcrRequest { + document: changed.document, + optional_params, + ..request + }) } pub(crate) struct PreparedOcrCall { diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 3e7507e9ed5..58ff841d83f 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -1,10 +1,9 @@ use std::future::Future; use std::pin::Pin; -use std::sync::Arc; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument}; +use super::types::{LiteLLMOcrResponse, OcrDocument}; use crate::Error; -use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; use serde::Serialize; use serde_json::Value; @@ -71,87 +70,3 @@ pub trait OcrHooks: Send + Sync { pub struct NoopOcrHooks; impl OcrHooks for NoopOcrHooks {} - -pub(crate) struct OcrLifecycleHooks { - pub hooks: Arc, - pub provider_name: String, -} - -impl CallLifecycleHooks - for OcrLifecycleHooks -{ - type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; - type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; - type SuccessFuture<'a> = OcrLogFuture<'a>; - type FailureFuture<'a> = OcrLogFuture<'a>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: LiteLLMOcrRequest, - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { - if !self.hooks.intercepts_requests() { - return Ok(request); - } - let changed = self - .hooks - .pre_call(OcrPreCallRequest { - model: request.model.clone(), - custom_llm_provider: self.provider_name.clone(), - document: request.document, - optional_params: Value::Object(request.optional_params), - }) - .await?; - let Value::Object(optional_params) = changed.optional_params else { - return Err(super::error::OcrRequestError::RequestField { - path: "guardrail.optional_params".into(), - } - .into()); - }; - Ok(LiteLLMOcrRequest { - document: changed.document, - optional_params, - ..request - }) - }) - } - - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: LiteLLMOcrRequest, - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { Ok(request) }) - } - - #[tracing::instrument( - name = "success_callback", - target = "litellm::function_trace", - level = "trace", - skip_all - )] - fn async_log_success_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - response: &'a LiteLLMOcrResponse, - timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - self.hooks.success(context, response, timing) - } - - #[tracing::instrument( - name = "failure_callback", - target = "litellm::function_trace", - level = "trace", - skip_all - )] - fn async_log_failure_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - error: &'a Error, - timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - self.hooks.failure(context, error, timing) - } -} diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index 92c9d4b717c..72fcda2811d 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -1,20 +1,21 @@ -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; - -use tokio::sync::{mpsc, oneshot}; +use std::time::{SystemTime, UNIX_EPOCH}; use super::handler::perform_ocr_request; use super::hooks::{ - OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, - OcrPreCallRequest, + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest, OcrPreCallRequest, }; use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; use crate::AuthError; use crate::Error; use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use crate::call_lifecycle::execution::HostExchange; use crate::call_lifecycle::host::{ - HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase, + HostCall, HostCallFuture, HostCallStep, HostFailure, HostPhase, LifecycleBackend, + LifecycleBackendFuture, +}; +use crate::call_lifecycle::workflow::{ + LifecycleCall, LifecycleOperation, Workflow, WorkflowFuture, WorkflowReply, }; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; @@ -93,15 +94,7 @@ pub enum OcrHostResult { pub type OcrCallStep = HostCallStep; -pub struct OcrCall { - lifecycle: HostLifecycle, - execution: OcrExecution, - response: Option>, - error: Option, - pending: bool, - completed: bool, - projecting: bool, -} +pub struct OcrCall(LifecycleCall); impl OcrCall { pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome { @@ -111,177 +104,21 @@ impl OcrCall { if !admission.host_operations { return NativeOutcome::Declined(OcrDecline::HostOperations); } - NativeOutcome::Completed(Self { - lifecycle: HostLifecycle::new(admission.asynchronous), - execution: OcrExecution::new(client), - response: None, - error: None, - pending: false, - completed: false, - projecting: false, - }) + NativeOutcome::Completed(Self(LifecycleCall::new( + OcrWorkflow { + client, + terminal: Arc::default(), + }, + admission.asynchronous, + ))) } pub async fn resume(&mut self, result: Option) -> Result { - if self.completed { - return Err(Error::InvalidRequest( - "OCR call cannot be resumed after completion".into(), - )); - } - if self.pending != result.is_some() { - return Err(Error::InvalidRequest( - "OCR host operation result does not match pending state".into(), - )); - } - match &result { - Some(OcrHostResult::Lifecycle(Ok(()))) - if self.lifecycle.phase() == HostPhase::Execute => - { - return Err(Error::InvalidRequest( - "OCR provider operation requires a typed result".into(), - )); - } - Some(result) - if !matches!(result, OcrHostResult::Lifecycle(_)) - && self.lifecycle.phase() != HostPhase::Execute => - { - return Err(Error::InvalidRequest( - "unexpected OCR provider operation result".into(), - )); - } - _ => {} - } - self.pending = false; - let provider_result = match result { - Some(OcrHostResult::Request(result)) if self.projecting => { - self.projecting = false; - match result { - Ok((request, azure_ad_token_provider)) => { - self.execution.request = Some(*request); - self.execution.azure_ad_token_provider = azure_ad_token_provider; - } - Err(error) => self.accept(Err(HostFailure::Error(error))), - } - None - } - Some(OcrHostResult::Request(_)) => { - return Err(Error::InvalidRequest( - "unexpected OCR request projection".into(), - )); - } - Some(OcrHostResult::Lifecycle(result)) => { - self.accept(result); - None - } - result => result, - }; - if self.lifecycle.phase() == HostPhase::Execute { - if self.execution.request.is_none() - && self.execution.execution.is_none() - && !self.execution.completed - { - self.projecting = true; - return Ok(self.host_step(OcrHostOperation::ProjectRequest)); - } - match self.execution.resume(provider_result).await { - Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)), - Ok(OcrCallStep::Complete(response)) => { - self.response = Some(Arc::new(response)); - self.accept(Ok(())); - } - Err(error) => self.accept(Err(HostFailure::Error(error))), - } - } - if self.error.is_some() { - self.execution.stop().await; - } - let operation = match self.lifecycle.phase() { - HostPhase::Complete => { - self.completed = true; - return match self.error.take() { - Some(error) => Err(error), - None => self - .response - .take() - .map(Arc::unwrap_or_clone) - .map(OcrCallStep::Complete) - .ok_or_else(|| { - Error::InvalidRequest("OCR completed without a response".into()) - }), - }; - } - HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse( - self.response - .as_ref() - .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? - .clone(), - ), - HostPhase::MapFailure => OcrHostOperation::MapFailure( - self.error - .as_ref() - .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? - .clone(), - ), - HostPhase::Success | HostPhase::Failure => { - let snapshot = self - .execution - .terminal - .lock() - .unwrap_or_else(|error| error.into_inner()) - .clone(); - match (self.lifecycle.phase(), snapshot) { - (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success { - context, - response: self - .response - .as_ref() - .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? - .clone(), - timing, - }, - (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure { - context, - error: self - .error - .as_ref() - .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? - .clone(), - timing, - }, - (phase, _) => OcrHostOperation::Lifecycle(phase), - } - } - phase => OcrHostOperation::Lifecycle(phase), - }; - Ok(self.host_step(operation)) - } - - fn accept(&mut self, result: Result<(), HostFailure>) { - let cancelled = matches!(&result, Err(HostFailure::Cancelled(_))); - if let Some(error) = self.lifecycle.accept(result) { - if cancelled { - self.error = Some(error); - } else { - self.error.get_or_insert(error); - } - self.execution.cancel(); - } + self.0.resume(result).await } pub async fn interrupt(&mut self, failure: HostFailure) -> Result { - if self.completed { - return Err(Error::InvalidRequest( - "OCR call cannot be interrupted after completion".into(), - )); - } - self.pending = false; - self.accept(Err(failure)); - self.resume(None).await - } - - fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep { - self.pending = true; - OcrCallStep::Host(operation) + self.0.interrupt(failure).await } } @@ -305,159 +142,148 @@ impl HostCall for OcrCall { } } -struct PendingOperation { - operation: OcrHostOperation, - result: oneshot::Sender, -} - -struct OcrExecution { - client: Option, - request: Option, - operations_tx: mpsc::UnboundedSender, - operations_rx: mpsc::UnboundedReceiver, - pending_result: Option>, - execution: Option>>, - completed: bool, - azure_ad_token_provider: bool, +struct OcrWorkflow { + client: OcrClient, terminal: Arc>>, } -impl OcrExecution { - fn new(client: OcrClient) -> Self { - let (operations_tx, operations_rx) = mpsc::unbounded_channel(); - Self { - client: Some(client), - request: None, - operations_tx, - operations_rx, - pending_result: None, - execution: None, - completed: false, - azure_ad_token_provider: false, - terminal: Arc::default(), +impl Workflow for OcrWorkflow { + type Request = (Box, bool); + type Operation = OcrHostOperation; + type Reply = OcrHostResult; + type Response = LiteLLMOcrResponse; + + fn operation(operation: LifecycleOperation) -> Self::Operation { + match operation { + LifecycleOperation::ProjectRequest => OcrHostOperation::ProjectRequest, + LifecycleOperation::Phase(phase) => OcrHostOperation::Lifecycle(phase), + LifecycleOperation::ConstructResponse(response) => { + OcrHostOperation::ConstructResponse(response) + } + LifecycleOperation::ConstructCachedResponse(_) => { + unreachable!("OCR does not cache responses") + } + LifecycleOperation::MapFailure(error) => OcrHostOperation::MapFailure(error), + LifecycleOperation::Success { + context, + response, + timing, + } => OcrHostOperation::Success { + context, + response, + timing, + }, + LifecycleOperation::Failure { + context, + error, + timing, + } => OcrHostOperation::Failure { + context, + error, + timing, + }, } } - pub async fn resume(&mut self, result: Option) -> Result { - if self.completed { - return Err(Error::InvalidRequest( - "OCR call cannot be resumed after completion".into(), - )); - } - match (self.pending_result.take(), result) { - (Some(sender), Some(result)) => sender - .send(result) - .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?, - (None, None) if self.execution.is_none() => self.start(), - (Some(sender), None) => { - self.pending_result = Some(sender); - return Err(Error::InvalidRequest( - "OCR host operation result is required".into(), - )); - } - (None, Some(_)) => { - return Err(Error::InvalidRequest( - "unexpected OCR host operation result".into(), - )); - } - (None, None) => {} - } + fn accepts(operation: &OcrHostOperation, reply: &OcrHostResult) -> bool { + matches!(reply, OcrHostResult::Lifecycle(Err(_))) + || matches!( + (operation, reply), + (OcrHostOperation::ProjectRequest, OcrHostResult::Request(_)) + | ( + OcrHostOperation::AcquireAzureAdToken, + OcrHostResult::AzureAdToken(_) + ) + | (OcrHostOperation::PreCall(_), OcrHostResult::PreCall(_)) + | ( + OcrHostOperation::DuringCall(_), + OcrHostResult::DuringCall(_) + ) + | (OcrHostOperation::PostCall(_), OcrHostResult::PostCall(_)) + | ( + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. }, + OcrHostResult::Lifecycle(_) + ) + ) + } - let execution = self.execution.as_mut().ok_or_else(|| { - Error::InvalidRequest("OCR call cannot be resumed after completion".into()) - })?; - tokio::select! { - operation = self.operations_rx.recv() => { - let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?; - self.pending_result = Some(operation.result); - Ok(OcrCallStep::Host(operation.operation)) - } - result = execution => { - self.execution = None; - self.completed = true; - result - .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))? - .map(OcrCallStep::Complete) - } + fn reply(&mut self, reply: OcrHostResult) -> WorkflowReply { + match reply { + OcrHostResult::Request(request) => WorkflowReply::Request(request), + OcrHostResult::Lifecycle(result) => WorkflowReply::Lifecycle(result), + reply => WorkflowReply::Operation(reply), } } - fn start(&mut self) { - let client = self.client.take().expect("admitted OCR call has a client"); - let mut request = self - .request - .take() - .expect("admitted OCR call has a request"); + fn start( + &mut self, + (request, azure_ad_token_provider): Self::Request, + operations: HostExchange, + ) -> WorkflowFuture { + let client = self.client.clone(); + let mut request = *request; let intercepts_requests = request.hooks.intercepts_requests(); - if self.azure_ad_token_provider { + let context = CallLifecycleContext::new( + "ocr", + request.model.clone(), + request.adapter.provider().as_str(), + request + .litellm_call_id + .clone() + .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::())), + ); + let started = epoch_seconds(); + if azure_ad_token_provider { request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new( OcrAzureAdTokenProvider { - operations: self.operations_tx.clone(), + operations: operations.clone(), }, ))); } request.hooks = Arc::new(ProtocolHooks { - operations: self.operations_tx.clone(), + operations, intercepts_requests, - terminal: self.terminal.clone(), }); - self.execution = Some(tokio::spawn(async move { - perform_ocr_request(&client, request).await - })); + let terminal = self.terminal.clone(); + Box::pin(async move { + let result = perform_ocr_request(&client, request).await; + let timing = CallLifecycleTiming::new(started, epoch_seconds()); + *terminal.lock().unwrap_or_else(|error| error.into_inner()) = Some((context, timing)); + result + }) } - fn cancel(&mut self) { - self.pending_result = None; - if let Some(execution) = &self.execution { - execution.abort(); - } - } - - async fn stop(&mut self) { - self.cancel(); - if let Some(execution) = self.execution.as_mut() { - let _ = execution.await; - } - self.execution = None; - } -} - -impl Drop for OcrExecution { - fn drop(&mut self) { - if let Some(execution) = &self.execution { - execution.abort(); - } + fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> { + self.terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() } } struct ProtocolHooks { - operations: mpsc::UnboundedSender, + operations: HostExchange, intercepts_requests: bool, - terminal: Arc>>, } #[derive(Debug)] struct OcrAzureAdTokenProvider { - operations: mpsc::UnboundedSender, + operations: HostExchange, } impl TokenProvider for OcrAzureAdTokenProvider { fn acquire(&self) -> TokenFuture<'_> { Box::pin(async move { - let (result, receiver) = oneshot::channel(); - self.operations - .send(PendingOperation { - operation: OcrHostOperation::AcquireAzureAdToken, - result, - }) - .map_err(|_| { - AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into()) - })?; - match receiver.await.map_err(|_| { - AuthError::AzureTokenAcquisition( - "OCR token provider operation was abandoned".into(), - ) - })? { + match self + .operations + .invoke(OcrHostOperation::AcquireAzureAdToken) + .await + .map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))? + { OcrHostResult::AzureAdToken(result) => result, _ => Err(AuthError::AzureTokenAcquisition( "invalid OCR token provider host result".into(), @@ -469,13 +295,7 @@ impl TokenProvider for OcrAzureAdTokenProvider { impl ProtocolHooks { async fn invoke(&self, operation: OcrHostOperation) -> Result { - let (result, receiver) = oneshot::channel(); - self.operations - .send(PendingOperation { operation, result }) - .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?; - receiver - .await - .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into())) + self.operations.invoke(operation).await } } @@ -519,48 +339,19 @@ impl OcrHooks for ProtocolHooks { } }) } - - fn success<'a>( - &'a self, - context: &'a CallLifecycleContext, - _response: &'a LiteLLMOcrResponse, - timing: &'a CallLifecycleTiming, - ) -> OcrLogFuture<'a> { - Box::pin(async move { - *self - .terminal - .lock() - .unwrap_or_else(|error| error.into_inner()) = - Some((context.clone(), timing.clone())); - }) - } - - fn failure<'a>( - &'a self, - context: &'a CallLifecycleContext, - _error: &'a Error, - timing: &'a CallLifecycleTiming, - ) -> OcrLogFuture<'a> { - Box::pin(async move { - *self - .terminal - .lock() - .unwrap_or_else(|error| error.into_inner()) = - Some((context.clone(), timing.clone())); - }) - } } -pub type OcrHostFuture<'a> = Pin + Send + 'a>>; - -pub trait OcrHost: Send + Sync { - fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>; +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) } pub struct NoopOcrHost; -impl OcrHost for NoopOcrHost { - fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { +impl LifecycleBackend for NoopOcrHost { + fn invoke(&self, operation: OcrHostOperation) -> LifecycleBackendFuture<'_, OcrHostResult> { Box::pin(async move { match operation { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( @@ -594,8 +385,8 @@ impl OcrHookHost { } } -impl OcrHost for OcrHookHost { - fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { +impl LifecycleBackend for OcrHookHost { + fn invoke(&self, operation: OcrHostOperation) -> LifecycleBackendFuture<'_, OcrHostResult> { Box::pin(async move { match operation { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 985844f47b0..259a50fb5b0 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -11,11 +11,12 @@ mod registry; pub mod types; pub mod wire; +pub use crate::call_lifecycle::host::LifecycleBackend as OcrHost; pub use client::{OcrClient, ocr}; pub use document::{encode_file_document, mime_type_for_name, upload_mime_type}; pub use lifecycle::{ NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, - OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, + OcrHookHost, OcrHostOperation, OcrHostResult, }; pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 6dc6b34b73d..65c22211d71 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -210,13 +210,18 @@ pub fn decode_request(wire: OcrWireRequest) -> Result } fn decode_document(value: Value) -> Result { + validate_document_url(&value)?; + decode_request_value(value, "document") +} + +pub fn validate_document_url(value: &Value) -> Result<(), OcrRequestError> { let kind = value.get("type").and_then(Value::as_str); let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); if missing_url { return Err(OcrRequestError::MissingDocumentUrl); } - decode_request_value(value, "document") + Ok(()) } fn source_for(sources: &BTreeMap, name: &str) -> InputSource { diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index e5e52bfce95..6c28b9de66f 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -1,9 +1,8 @@ use std::collections::BTreeMap; -use std::sync::{Mutex, OnceLock}; +use std::sync::OnceLock; use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; -use crate::caching::in_memory_cache::InMemoryCache; use crate::error::Error; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; @@ -12,6 +11,7 @@ use aws_sigv4::http_request::{ }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use litellm_cache_memory::InMemoryCache; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; @@ -26,7 +26,7 @@ use super::constants::{ const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); -static IAM_CREDENTIALS_CACHE: OnceLock>> = OnceLock::new(); +static IAM_CREDENTIALS_CACHE: OnceLock> = OnceLock::new(); fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { match flow { @@ -108,16 +108,17 @@ fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { } fn get_cached_credentials(key: &str) -> Option { - let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); - let mut entries = cache.lock().ok()?; - entries.get_cache(key) + IAM_CREDENTIALS_CACHE + .get_or_init(InMemoryCache::default) + .get_cache(key) + .ok() + .flatten() } fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) { - let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); - if let Ok(mut entries) = cache.lock() { - entries.set_cache(key, credentials, Some(ttl)); - } + let _ = IAM_CREDENTIALS_CACHE + .get_or_init(InMemoryCache::default) + .set_cache(key, credentials, Some(ttl)); } fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index b1098f4d386..7414a830314 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -1,12 +1,8 @@ -use std::future::Future; -use std::pin::Pin; use std::sync::Mutex; use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::Value; -use crate::Error; -use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType}; #[derive(Clone, Debug, Default, PartialEq, Eq)] @@ -203,58 +199,16 @@ impl ResponsesWsInstrumentation { } }) } -} -type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; - -impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { - type PreCallFuture<'a> = LifecycleFuture<'a, ()>; - type DuringCallFuture<'a> = LifecycleFuture<'a, ()>; - type SuccessFuture<'a> = Pin + Send + 'a>>; - type FailureFuture<'a> = Pin + Send + 'a>>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: (), - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { Ok(request) }) - } - - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: (), - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { Ok(request) }) - } - - fn async_log_success_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _response: &'a (), - _timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - Box::pin(async move { - let outcome = self.success_outcome(); - if let Ok(mut state) = self.state.lock() { - state.outcome = Some(outcome); - } - }) - } - - fn async_log_failure_event<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _error: &'a Error, - _timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - Box::pin(async move { - let outcome = self.failure_outcome(); - if let Ok(mut state) = self.state.lock() { - state.outcome = Some(outcome); - } - }) + pub fn record_outcome(&self, success: bool) { + let outcome = if success { + self.success_outcome() + } else { + self.failure_outcome() + }; + if let Ok(mut state) = self.state.lock() { + state.outcome = Some(outcome); + } } } @@ -328,25 +282,12 @@ mod tests { )); } - #[tokio::test] - async fn lifecycle_records_success_outcome_for_provider_completion() { + #[test] + fn records_success_outcome_for_provider_completion() { let instrumentation = ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default()); - let result = crate::call_lifecycle::CallLifecycle::default() - .run( - crate::call_lifecycle::CallLifecycleContext::new( - "responses_websocket", - "gpt-5", - "openai", - "call-1", - ), - (), - &instrumentation, - |_| async { Ok::<(), Error>(()) }, - ) - .await; + instrumentation.record_outcome(true); - assert!(result.is_ok()); assert!(matches!( instrumentation.take_outcome(), Some(ResponsesWsLogOutcome::Success { .. }) diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs index 19fb946afde..5424809bdc0 100644 --- a/litellm-rust/crates/core/tests/host_lifecycle.rs +++ b/litellm-rust/crates/core/tests/host_lifecycle.rs @@ -27,6 +27,34 @@ fn public_outcome_is_finalized_before_a_single_terminal_dispatch() { for asynchronous in [false, true] { let (events, failures) = run(None, asynchronous); assert!(failures.is_empty()); + let expected = if asynchronous { + vec![ + HostPhase::Setup, + HostPhase::DeploymentPreCall, + HostPhase::Prepare, + HostPhase::CacheLookup, + HostPhase::Execute, + HostPhase::ConstructResponse, + HostPhase::PostProcess, + HostPhase::DeploymentPostCall, + HostPhase::CacheStore, + HostPhase::Finalize, + HostPhase::Success, + ] + } else { + vec![ + HostPhase::Setup, + HostPhase::Prepare, + HostPhase::CacheLookup, + HostPhase::Execute, + HostPhase::ConstructResponse, + HostPhase::PostProcess, + HostPhase::CacheStore, + HostPhase::Finalize, + HostPhase::Success, + ] + }; + assert_eq!(events, expected); assert_eq!( &events[events.len() - 2..], &[HostPhase::Finalize, HostPhase::Success] @@ -51,9 +79,12 @@ fn only_provider_and_response_construction_failures_use_provider_mapping() { HostPhase::Setup, HostPhase::DeploymentPreCall, HostPhase::Prepare, + HostPhase::CacheLookup, HostPhase::Execute, HostPhase::ConstructResponse, + HostPhase::PostProcess, HostPhase::DeploymentPostCall, + HostPhase::CacheStore, HostPhase::Finalize, ] { let (events, failures) = run(Some(phase), true); @@ -76,6 +107,35 @@ fn only_provider_and_response_construction_failures_use_provider_mapping() { } } +#[test] +fn cache_hit_uses_the_same_graph_without_provider_or_cache_store() { + let mut lifecycle = HostLifecycle::new(true); + let mut events = Vec::new(); + while lifecycle.phase() != HostPhase::CacheLookup { + events.push(lifecycle.phase()); + lifecycle.accept(Ok(())); + } + events.push(HostPhase::CacheLookup); + lifecycle.cache_hit(); + while lifecycle.phase() != HostPhase::Complete { + events.push(lifecycle.phase()); + lifecycle.accept(Ok(())); + } + assert_eq!( + events, + [ + HostPhase::Setup, + HostPhase::DeploymentPreCall, + HostPhase::Prepare, + HostPhase::CacheLookup, + HostPhase::ConstructResponse, + HostPhase::PostProcess, + HostPhase::Finalize, + HostPhase::Success, + ] + ); +} + #[test] fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { let mut lifecycle = HostLifecycle::new(true); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 55f8713d76e..22c3e5bc10a 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -490,7 +490,9 @@ async fn direct_native_host_drives_the_same_state_machine() { "DuringCall", "PostCall", "ConstructResponse", + "PostProcess", "DeploymentPostCall", + "CacheStore", "Finalize", "Success", ] diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 42fad740870..57460566f64 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,8 @@ trace-parity = [ futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-cache.workspace = true +litellm-cache-memory.workspace = true litellm-token-counter.workspace = true litellm-python-interop.workspace = true pyo3.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 90ba5b65b6f..51087403660 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,4 +1,5 @@ mod auth; +mod cache; mod constants; mod diagnostics; mod errors; diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs index 586c23e0ad4..1b1f5c6cb87 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -3,6 +3,8 @@ use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; +use super::contract::CallbackPhase; + #[derive(FromPyObject)] pub(crate) struct PythonLogger(Py); @@ -19,7 +21,7 @@ impl PythonLogger { visit.call(&self.0) } - pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { + pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: CallbackPhase) -> PyResult { if !self .object(py) .getattr("_native_callback_fast_path") @@ -29,7 +31,7 @@ impl PythonLogger { } py.import("litellm.rust_bridge.lifecycle")? .getattr("callbacks_needed")? - .call1((self.object(py), phase))? + .call1((self.object(py), phase.as_str()))? .extract() } @@ -68,7 +70,7 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { + if !self.callbacks_needed(py, CallbackPhase::SyncSuccessForAsyncCall)? { return Ok(()); } self.object(py).call_method1( @@ -89,9 +91,9 @@ impl PythonLogger { if !self.callbacks_needed( py, if asynchronous { - "async_failure" + CallbackPhase::AsyncFailure } else { - "sync_failure" + CallbackPhase::SyncFailure }, )? { py.import("litellm.rust_bridge.lifecycle")? @@ -129,7 +131,7 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { + if !self.callbacks_needed(py, CallbackPhase::SyncSuccess)? { return self.success_bookkeeping(py, response, start, end, false); } let context = py.import("contextvars")?.call_method0("copy_context")?; @@ -155,7 +157,7 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { + if !self.callbacks_needed(py, CallbackPhase::AsyncSuccess)? { return self.success_bookkeeping(py, response, start, end, true); } let context = py.import("contextvars")?.call_method0("copy_context")?; diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/completed.rs b/litellm-rust/crates/python-bridge/src/lifecycle/completed.rs new file mode 100644 index 00000000000..96fe644a49f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/completed.rs @@ -0,0 +1,281 @@ +use std::marker::PhantomData; + +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use serde::Serialize; + +use litellm_core::call_lifecycle::host::{HostPhase, HostStep}; +use litellm_core::call_lifecycle::provider::{ + CompletedCall, CompletedOperation, CompletedReply, CompletedRoute, CompletedWorkflow, + ProviderRequest, ProviderResponse, +}; +use litellm_core::call_lifecycle::workflow::LifecycleOperation; +use litellm_python_interop::{ + from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +}; + +use super::contract::{AdapterOperation, CallMode, PythonCallType}; +use super::{OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call}; +use crate::errors::execution_error_to_pyerr; + +pub(crate) trait PythonCompletedRoute: CompletedRoute { + const SYNC_CALL_TYPE: PythonCallType; + const ASYNC_CALL_TYPE: PythonCallType; + + fn admit(request: &Bound<'_, PyDict>) -> PyResult<()>; + fn project(request: &Bound<'_, PyDict>) -> PyResult; +} + +struct PythonCompletedHost { + state: PythonCallState, + request: Py, + adapter: Py, + pending: Option, + route: PhantomData, +} + +impl PythonRoute for PythonCompletedHost +where + R::Response: Serialize, +{ + type Call = CompletedCall; + + fn state(&self) -> &PythonCallState { + &self.state + } + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.state + } + + fn classify(operation: &CompletedOperation) -> OperationClass { + match operation { + CompletedOperation::Lifecycle(LifecycleOperation::Phase( + HostPhase::Prepare | HostPhase::PostProcess | HostPhase::CacheStore, + )) => OperationClass::Route, + CompletedOperation::Lifecycle(LifecycleOperation::Phase(phase)) => { + OperationClass::Phase(*phase) + } + CompletedOperation::Lifecycle(LifecycleOperation::Success { .. }) => { + OperationClass::Phase(HostPhase::Success) + } + CompletedOperation::Lifecycle(LifecycleOperation::Failure { .. }) => { + OperationClass::Phase(HostPhase::Failure) + } + _ => OperationClass::Route, + } + } + + fn lifecycle_result() -> CompletedReply { + CompletedReply::Lifecycle(Ok(())) + } + fn map_error(error: litellm_core::Error) -> PyErr { + execution_error_to_pyerr(error) + } + + fn invoke( + &mut self, + py: Python<'_>, + operation: CompletedOperation, + ) -> PyResult> { + match self.invoke_step(py, operation)? { + HostStep::Ready(reply) => Ok(reply), + HostStep::Suspend(_) => Err(pyo3::exceptions::PyRuntimeError::new_err( + "route operation requires async execution", + )), + } + } + + fn invoke_step( + &mut self, + py: Python<'_>, + operation: CompletedOperation, + ) -> PyResult, Py>> { + let (operation, payload) = match operation { + CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::Prepare)) => { + self.state.prepare(py)?; + return Ok(HostStep::Ready(CompletedReply::Prepared(Ok( + crate::cache::plan( + py, + self.state.call_type.as_str(), + self.state.kwargs.bind(py), + )?, + )))); + } + CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::PostProcess)) => ( + AdapterOperation::PostProcess, + self.state + .response + .as_ref() + .ok_or_else(missing_state)? + .clone_ref(py), + ), + CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::CacheStore)) => ( + AdapterOperation::CacheStore, + self.state + .response + .as_ref() + .ok_or_else(missing_state)? + .clone_ref(py), + ), + CompletedOperation::Lifecycle(LifecycleOperation::ConstructCachedResponse( + response, + )) => { + self.state.end = Some(now(py)?); + (AdapterOperation::CachedResponse, to_py(py, &response)?) + } + CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => ( + AdapterOperation::ProjectRequest, + self.request.clone_ref(py).into_any(), + ), + CompletedOperation::BeforeRequest(request) => { + (AdapterOperation::BeforeRequest, to_py(py, &request)?) + } + CompletedOperation::AfterResponse(response) => { + (AdapterOperation::AfterResponse, to_py(py, &response)?) + } + CompletedOperation::Lifecycle(LifecycleOperation::ConstructResponse(response)) => { + self.state.end = Some(now(py)?); + ( + AdapterOperation::ConstructResponse, + to_py(py, response.as_ref())?, + ) + } + CompletedOperation::Lifecycle(LifecycleOperation::MapFailure(error)) => { + if self.state.error.is_none() { + self.state.retain_error(py, execution_error_to_pyerr(error)); + } + if self.state.end.is_none() { + self.state.end = Some(now(py)?); + } + ( + AdapterOperation::MapFailure, + self.state + .error + .as_ref() + .ok_or_else(missing_state)? + .clone_ref(py) + .into_any(), + ) + } + _ => return Err(missing_state()), + }; + self.pending = Some(operation); + let step = self.adapter.bind(py).call_method1( + "invoke", + ( + operation.as_str(), + payload, + &self.request, + &self.state.kwargs, + self.state.logger()?.object(py), + ), + )?; + let protocol = py.import("litellm.rust_bridge.lifecycle")?; + if step.is_instance(&protocol.getattr("Await")?)? { + if !self.state.mode.is_async() { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "sync route operation suspended", + )); + } + return Ok(HostStep::Suspend(step.getattr("awaitable")?.unbind())); + } + if !step.is_instance(&protocol.getattr("Complete")?)? { + return Err(missing_state()); + } + self.accept_route(py, step.getattr("value")?.unbind()) + .map(HostStep::Ready) + } + + fn accept_route( + &mut self, + py: Python<'_>, + value: Py, + ) -> PyResult> { + Ok(match self.pending.take().ok_or_else(missing_state)? { + AdapterOperation::PostProcess => Self::lifecycle_result(), + AdapterOperation::CacheStore => CompletedReply::CacheStore(Ok(if value.is_none(py) { + None + } else { + Some(from_py(value.bind(py))?) + })), + AdapterOperation::ProjectRequest => { + let request = value.into_bound(py).cast_into::()?; + CompletedReply::Request(Ok(R::project(&request)?)) + } + AdapterOperation::BeforeRequest => { + CompletedReply::BeforeRequest(Ok(from_py::(value.bind(py))?)) + } + AdapterOperation::AfterResponse => { + CompletedReply::AfterResponse(Ok(from_py::(value.bind(py))?)) + } + AdapterOperation::CachedResponse | AdapterOperation::ConstructResponse => { + self.state.response = Some(value); + Self::lifecycle_result() + } + AdapterOperation::MapFailure => { + self.state + .retain_error(py, PyErr::from_value(value.into_bound(py))); + Self::lifecycle_result() + } + }) + } + + fn cleanup(&mut self) { + self.pending = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.request)?; + visit.call(&self.adapter) + } +} + +pub(crate) fn run( + py: Python<'_>, + request: Bound<'_, PyDict>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, + host: Bound<'_, PyAny>, +) -> PyResult> +where + R::Response: Serialize, +{ + R::admit(&request)?; + let controls = crate::cache::snapshot( + py, + if asynchronous { + R::ASYNC_CALL_TYPE.as_str() + } else { + R::SYNC_CALL_TYPE.as_str() + }, + &request, + )?; + crate::errors::admit( + litellm_core::call_lifecycle::cache::ResponseCachePlan { + controls, + ..Default::default() + } + .admit(), + )?; + let call = CompletedCall::::new(CompletedWorkflow::default(), asynchronous); + let host = PythonCompletedHost:: { + state: PythonCallState::new( + py, + args.unbind(), + kwargs.copy()?.unbind(), + CallMode::from_async(asynchronous), + if asynchronous { + R::ASYNC_CALL_TYPE + } else { + R::SYNC_CALL_TYPE + }, + )?, + request: request.unbind(), + adapter: host.unbind(), + pending: None, + route: PhantomData, + }; + run_call(py, call, host) +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/contract.rs b/litellm-rust/crates/python-bridge/src/lifecycle/contract.rs new file mode 100644 index 00000000000..d9028223eeb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/contract.rs @@ -0,0 +1,163 @@ +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyString; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum CallMode { + Sync, + Async, +} + +impl CallMode { + pub(crate) fn from_async(asynchronous: bool) -> Self { + if asynchronous { + Self::Async + } else { + Self::Sync + } + } + + pub(crate) fn is_async(self) -> bool { + matches!(self, Self::Async) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PythonCallType { + Ocr, + AsyncOcr, + Completion, + AsyncCompletion, + AnthropicMessages, + Transcription, + AsyncTranscription, + #[cfg(test)] + Synthetic, + #[cfg(test)] + Test, +} + +impl PythonCallType { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Ocr => "ocr", + Self::AsyncOcr => "aocr", + Self::Completion => "completion", + Self::AsyncCompletion => "acompletion", + Self::AnthropicMessages => "anthropic_messages", + Self::Transcription => "transcription", + Self::AsyncTranscription => "atranscription", + #[cfg(test)] + Self::Synthetic => "synthetic", + #[cfg(test)] + Self::Test => "test", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum CallbackPhase { + Input, + Payload, + SyncSuccess, + SyncSuccessForAsyncCall, + AsyncSuccess, + SyncFailure, + AsyncFailure, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum AdapterOperation { + PostProcess, + CacheStore, + CachedResponse, + ProjectRequest, + BeforeRequest, + AfterResponse, + ConstructResponse, + MapFailure, +} + +impl AdapterOperation { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::PostProcess => "post_process", + Self::CacheStore => "cache_response", + Self::CachedResponse => "cached_response", + Self::ProjectRequest => "project", + Self::BeforeRequest => "before_request", + Self::AfterResponse => "after_response", + Self::ConstructResponse => "response", + Self::MapFailure => "map_failure", + } + } +} + +impl CallbackPhase { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Input => "input", + Self::Payload => "payload", + Self::SyncSuccess => "sync_success", + Self::SyncSuccessForAsyncCall => "sync_success_async", + Self::AsyncSuccess => "async_success", + Self::SyncFailure => "sync_failure", + Self::AsyncFailure => "async_failure", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum RequestField { + ApiBase, + ApiKey, + Audio, + Body, + CustomLlmProvider, + ExtraHeaders, + HasAgenticHook, + HostFacts, + LitellmCallId, + Messages, + Model, + OptionalParams, + Timeout, +} + +impl RequestField { + pub(crate) fn key(self, py: Python<'_>) -> &Bound<'_, PyString> { + match self { + Self::ApiBase => intern!(py, "api_base"), + Self::ApiKey => intern!(py, "api_key"), + Self::Audio => intern!(py, "audio"), + Self::Body => intern!(py, "body"), + Self::CustomLlmProvider => intern!(py, "custom_llm_provider"), + Self::ExtraHeaders => intern!(py, "extra_headers"), + Self::HasAgenticHook => intern!(py, "has_agentic_hook"), + Self::HostFacts => intern!(py, "host_facts"), + Self::LitellmCallId => intern!(py, "litellm_call_id"), + Self::Messages => intern!(py, "messages"), + Self::Model => intern!(py, "model"), + Self::OptionalParams => intern!(py, "optional_params"), + Self::Timeout => intern!(py, "timeout"), + } + } + + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::ApiBase => "api_base", + Self::ApiKey => "api_key", + Self::Audio => "audio", + Self::Body => "body", + Self::CustomLlmProvider => "custom_llm_provider", + Self::ExtraHeaders => "extra_headers", + Self::HasAgenticHook => "has_agentic_hook", + Self::HostFacts => "host_facts", + Self::LitellmCallId => "litellm_call_id", + Self::Messages => "messages", + Self::Model => "model", + Self::OptionalParams => "optional_params", + Self::Timeout => "timeout", + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs index a5579d0f2ae..31a4b022d4c 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs @@ -6,6 +6,7 @@ use pyo3::exceptions::PyException; use pyo3::prelude::*; use super::bindings::PythonLogger; +use super::contract::CallbackPhase; use super::state::PythonCallState; impl PythonCallState { @@ -32,7 +33,7 @@ impl PythonCallState { state: self, logger, }; - match success_dispatch(self.asynchronous, self.internal, &facts)? { + match success_dispatch(self.mode.is_async(), self.internal, &facts)? { SuccessDispatch::SyncBookkeeping => { logger.success_bookkeeping(py, &self.response, &self.start, &self.end, false) } @@ -62,7 +63,7 @@ impl PythonCallState { py: Python<'_>, asynchronous: bool, ) -> PyResult>> { - if !failure_dispatch(self.asynchronous, self.internal, self.logger.is_some()) { + if !failure_dispatch(self.mode.is_async(), self.internal, self.logger.is_some()) { return Ok(None); } let Some(error) = &self.error else { @@ -95,9 +96,9 @@ impl SuccessFacts for PythonSuccessFacts<'_, '_> { self.logger.callbacks_needed( self.py, if asynchronous { - "async_success" + CallbackPhase::AsyncSuccess } else { - "sync_success" + CallbackPhase::SyncSuccess }, ) } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs index 17a480a7225..919868187fc 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs @@ -1,10 +1,11 @@ use std::panic::{AssertUnwindSafe, catch_unwind}; -use litellm_python_interop::panic_to_pyerr; use pyo3::exceptions::{PyBaseException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; +use litellm_python_interop::panic_to_pyerr; + pub(super) enum ExecutionStep { Return(Py), Await(Py), diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 7af600bb5ed..bb61182a83d 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -1,7 +1,10 @@ mod bindings; +pub(crate) mod completed; +pub(crate) mod contract; mod dispatch; mod handle; mod preparation; +pub(crate) mod request; mod runner; mod state; diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/request.rs b/litellm-rust/crates/python-bridge/src/lifecycle/request.rs new file mode 100644 index 00000000000..2013b7de541 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/request.rs @@ -0,0 +1,63 @@ +use litellm_core::call_lifecycle::provider::ProviderOptions; +use litellm_python_interop::from_py_preserving_errors as from_py; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Map, Value}; + +use super::contract::RequestField; + +pub(crate) fn required<'py>( + request: &Bound<'py, PyDict>, + field: RequestField, +) -> PyResult> { + let name = field.as_str(); + request + .get_item(field.key(request.py()))? + .ok_or_else(|| pyo3::exceptions::PyTypeError::new_err(format!("missing {name}"))) +} + +pub(crate) fn optional_string( + request: &Bound<'_, PyDict>, + field: RequestField, +) -> PyResult> { + request + .get_item(field.key(request.py()))? + .map(|value| value.extract()) + .transpose() + .map(Option::flatten) +} + +pub(crate) fn object( + request: &Bound<'_, PyDict>, + field: RequestField, +) -> PyResult> { + request + .get_item(field.key(request.py()))? + .filter(|value| !value.is_none()) + .map(|value| from_py(&value)) + .transpose() + .map(Option::unwrap_or_default) +} + +pub(crate) fn options(request: &Bound<'_, PyDict>) -> PyResult { + Ok(ProviderOptions { + model: required(request, RequestField::Model)?.extract()?, + litellm_call_id: optional_string(request, RequestField::LitellmCallId)?, + api_key: optional_string(request, RequestField::ApiKey)?, + api_base: optional_string(request, RequestField::ApiBase)?, + custom_llm_provider: optional_string(request, RequestField::CustomLlmProvider)?, + extra_headers: request + .get_item(RequestField::ExtraHeaders.key(request.py()))? + .filter(|value| !value.is_none()) + .map(|value| from_py(&value)) + .transpose()?, + timeout: crate::marshal::optional_timeout( + request + .get_item(RequestField::Timeout.key(request.py()))? + .filter(|value| !value.is_none()) + .map(|value| crate::marshal::python_timeout_seconds(request.py(), value.unbind())) + .transpose()? + .flatten(), + ), + }) +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs b/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs index 26760f743e5..8efd7953c01 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs @@ -2,13 +2,15 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_core::call_lifecycle::host::{ - HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, -}; +use tokio::sync::Mutex; + use pyo3::exceptions::{PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; -use tokio::sync::Mutex; + +use litellm_core::call_lifecycle::host::{ + HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, +}; use super::handle::{Execution, ExecutionBody, ExecutionStep}; use super::state::{PythonCallState, missing_state, now}; @@ -32,6 +34,20 @@ pub(crate) trait PythonRoute: Send + Sync { py: Python<'_>, operation: ::Operation, ) -> PyResult<::Result>; + fn invoke_step( + &mut self, + py: Python<'_>, + operation: ::Operation, + ) -> PyResult::Result, Py>> { + self.invoke(py, operation).map(HostStep::Ready) + } + fn accept_route( + &mut self, + _py: Python<'_>, + _value: Py, + ) -> PyResult<::Result> { + Err(missing_state()) + } fn cleanup(&mut self); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } @@ -48,6 +64,7 @@ struct NativeCallState { enum PendingOperation { Native, Host(HostPhase), + Route, } struct PythonLifecycle { @@ -62,7 +79,7 @@ pub(crate) fn run_call( call: R::Call, route: R, ) -> PyResult> { - let asynchronous = route.state().asynchronous; + let asynchronous = route.state().mode.is_async(); let mut lifecycle = PythonLifecycle { route, call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))), @@ -102,7 +119,7 @@ impl PythonLifecycle { call.result = Some(result); Ok(()) }; - if self.route.state().asynchronous { + if self.route.state().mode.is_async() { let mut future = Box::pin(future); if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? { return Ok(HostStep::Ready(self.take_native_result()?)); @@ -182,6 +199,11 @@ impl PythonLifecycle { }; self.resume_core(py, Some(result))? } + (Some(PendingOperation::Route), Some(result)) => { + let result = result.and_then(|value| self.route.accept_route(py, value)); + let result = result.map_err(|error| self.host_failure(py, error, None)); + self.resume_core(py, Some(result))? + } _ => return Err(missing_state()), }; loop { @@ -215,7 +237,14 @@ impl PythonLifecycle { .map(|()| R::lifecycle_result()), Err(error) => Err(error), }, - None => self.route.invoke(py, operation), + None => match self.route.invoke_step(py, operation) { + Ok(HostStep::Ready(result)) => Ok(result), + Ok(HostStep::Suspend(awaitable)) => { + self.pending = Some(PendingOperation::Route); + return Ok(ExecutionStep::Await(awaitable)); + } + Err(error) => Err(error), + }, }; let result = match result { Ok(result) => Ok(result), diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/state.rs b/litellm-rust/crates/python-bridge/src/lifecycle/state.rs index 3c163123962..908a8c88786 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/state.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/state.rs @@ -5,6 +5,7 @@ use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; use super::bindings::{self, DeploymentHooks, PythonLogger}; +use super::contract::{CallMode, PythonCallType}; use super::preparation; pub(crate) fn missing_state() -> PyErr { @@ -19,9 +20,9 @@ pub(crate) struct PythonCallState { pub end: Option>, pub response: Option>, pub error: Option>, - pub asynchronous: bool, + pub mode: CallMode, pub internal: bool, - pub call_type: &'static str, + pub call_type: PythonCallType, } pub(crate) fn now(py: Python<'_>) -> PyResult> { @@ -46,7 +47,7 @@ impl PythonCallState { return Ok(HostStep::Suspend(DeploymentHooks::before_call( py, &self.kwargs, - self.call_type, + self.call_type.as_str(), )?)); } HostPhase::Prepare => self.prepare(py)?, @@ -62,7 +63,7 @@ impl PythonCallState { py, &self.kwargs, &self.response, - self.call_type, + self.call_type.as_str(), )?)); } HostPhase::Finalize => self.finalize(py)?, @@ -75,7 +76,7 @@ impl PythonCallState { py, &self.kwargs, error, - self.call_type, + self.call_type.as_str(), )?)); } } @@ -86,7 +87,9 @@ impl PythonCallState { return Ok(HostStep::Suspend(awaitable)); } } + HostPhase::PostProcess | HostPhase::CacheStore => {} HostPhase::Execute + | HostPhase::CacheLookup | HostPhase::ConstructResponse | HostPhase::MapFailure | HostPhase::Complete => return Err(missing_state()), @@ -114,8 +117,8 @@ impl PythonCallState { py: Python<'_>, args: Py, kwargs: Py, - asynchronous: bool, - call_type: &'static str, + mode: CallMode, + call_type: PythonCallType, ) -> PyResult { Ok(Self { args, @@ -125,7 +128,7 @@ impl PythonCallState { end: None, response: None, error: None, - asynchronous, + mode, internal: false, call_type, }) @@ -142,11 +145,11 @@ impl PythonCallState { self.internal = bindings::is_internal_call(py)?; let result = bindings::setup( py, - self.call_type, + self.call_type.as_str(), &self.args, &self.kwargs, &self.start, - self.asynchronous, + self.mode.is_async(), )?; self.logger = Some(result.logger()?); self.kwargs = result.kwargs()?; diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs b/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs index 0b0ee88896d..f1320f33c54 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs @@ -6,6 +6,7 @@ use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyTuple; +use super::contract::{CallMode, PythonCallType}; use super::dispatch::{PendingLogging, PendingSuccess}; use super::handle::{Execution, ExecutionBody, ExecutionStep}; use super::*; @@ -160,8 +161,8 @@ fn shared_runner_executes_a_non_ocr_adapter() { py, PyTuple::empty(py).unbind(), PyDict::new(py).unbind(), - false, - "synthetic", + CallMode::Sync, + PythonCallType::Synthetic, ) .unwrap(), ); @@ -196,8 +197,8 @@ fn ready_native_lifecycle_completes_without_scheduling() { py, PyTuple::empty(py).unbind(), PyDict::new(py).unbind(), - true, - "synthetic", + CallMode::Async, + PythonCallType::Synthetic, ) .unwrap(), ); @@ -280,8 +281,8 @@ fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Executi py, PyTuple::empty(py).unbind(), PyDict::new(py).unbind(), - true, - "test", + CallMode::Async, + PythonCallType::Test, ) .unwrap(); state.retain_error(py, PyErr::from_value(error.into_any())); @@ -342,9 +343,9 @@ fn state( end: Some(py.None()), response: Some(response), error: None, - asynchronous, + mode: CallMode::from_async(asynchronous), internal: false, - call_type: "test", + call_type: PythonCallType::Test, } } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/lifecycle.rs index 99d29258076..fde08844ae0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/lifecycle.rs @@ -1,2 +1,57 @@ -// TODO: implement chat_completions lifecycle checkpoints before replacing the Python lifecycle -unimplemented_lifecycle_route!(ChatCompletions, _chat_completions_lifecycle); +use litellm_core::chat_completions::lifecycle::{ + ChatCompletionsRoute, OwnedChatCompletionsRequest, +}; +use litellm_python_interop::from_py_preserving_errors as from_py; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use crate::lifecycle::completed::{self, PythonCompletedRoute}; +use crate::lifecycle::contract::{PythonCallType, RequestField}; +use crate::lifecycle::request::{object, optional_string, options, required}; + +impl PythonCompletedRoute for ChatCompletionsRoute { + const SYNC_CALL_TYPE: PythonCallType = PythonCallType::Completion; + const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AsyncCompletion; + + fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> { + crate::errors::admit(litellm_core::chat_completions::admit( + &required(request, RequestField::Model)?.extract::()?, + optional_string(request, RequestField::CustomLlmProvider)?.as_deref(), + from_py(&required(request, RequestField::Messages)?)?, + &object(request, RequestField::OptionalParams)?, + Some(&object(request, RequestField::ExtraHeaders)?), + request + .get_item(RequestField::HostFacts.key(request.py()))? + .map(|value| from_py(&value)) + .transpose()? + .unwrap_or_default(), + )) + } + + fn project(request: &Bound<'_, PyDict>) -> PyResult { + Ok(OwnedChatCompletionsRequest { + options: options(request)?, + messages: from_py(&required(request, RequestField::Messages)?)?, + optional_params: object(request, RequestField::OptionalParams)?, + }) + } +} + +#[pyfunction] +fn _chat_completions_lifecycle( + py: Python<'_>, + request: Bound<'_, PyDict>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, + host: Bound<'_, PyAny>, +) -> PyResult> { + completed::run::(py, request, args, kwargs, asynchronous, host) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + crate::routes::definition::add_function( + module, + wrap_pyfunction!(_chat_completions_lifecycle, module)?, + ) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs index c40935da576..711be79496e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs @@ -88,7 +88,7 @@ bridge_route! { extra_headers: Option, timeout_seconds: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - host_facts: Option, + host_facts: Option, on_request: Option>, }, prepare = prepare_chat_completions, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/messages/lifecycle.rs index ac1bfd06c1d..8f12606bb75 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/lifecycle.rs @@ -1,2 +1,49 @@ -// TODO: implement messages lifecycle checkpoints before replacing the Python lifecycle -unimplemented_lifecycle_route!(Messages, _messages_lifecycle); +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use litellm_core::messages::lifecycle::{MessagesRoute, OwnedMessagesRequest}; +use litellm_python_interop::from_py_preserving_errors as from_py; + +use crate::lifecycle::completed::{self, PythonCompletedRoute}; +use crate::lifecycle::contract::{PythonCallType, RequestField}; +use crate::lifecycle::request::{optional_string, options, required}; + +impl PythonCompletedRoute for MessagesRoute { + const SYNC_CALL_TYPE: PythonCallType = PythonCallType::AnthropicMessages; + const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AnthropicMessages; + + fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> { + crate::errors::admit(litellm_core::messages::admit( + &required(request, RequestField::Model)?.extract::()?, + optional_string(request, RequestField::CustomLlmProvider)?.as_deref(), + request + .get_item(RequestField::HasAgenticHook.key(request.py()))? + .map(|value| value.extract()) + .transpose()? + .unwrap_or(false), + )) + } + + fn project(request: &Bound<'_, PyDict>) -> PyResult { + Ok(OwnedMessagesRequest { + options: options(request)?, + body: from_py(&required(request, RequestField::Body)?)?, + }) + } +} + +#[pyfunction] +fn _messages_lifecycle( + py: Python<'_>, + request: Bound<'_, PyDict>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, + host: Bound<'_, PyAny>, +) -> PyResult> { + completed::run::(py, request, args, kwargs, asynchronous, host) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + crate::routes::definition::add_function(module, wrap_pyfunction!(_messages_lifecycle, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index 034c7b1d681..5cf1b5bbb99 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -8,6 +8,7 @@ use litellm_core::ocr::hooks::OcrPreCallRequest; use litellm_python_interop::to_py_preserving_errors as to_py; use crate::lifecycle::PythonLogger; +use crate::lifecycle::contract::CallbackPhase; pub(super) struct OcrLoggingFields { model: String, @@ -89,7 +90,7 @@ impl PythonLogger { kwargs.set_item("input", "OCR document processing")?; kwargs.set_item("api_key", api_key)?; kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { + if self.callbacks_needed(py, CallbackPhase::Input)? { self.object(py).call_method("pre_call", (), Some(&kwargs))?; } else { self.object(py) @@ -109,7 +110,7 @@ impl PythonLogger { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { + if self.callbacks_needed(py, CallbackPhase::Input)? { let kwargs = PyDict::new(py); kwargs.set_item("original_response", to_py(py, original_response)?)?; kwargs.set_item("additional_args", &additional)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index f5ede86a64b..beff6da7027 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -11,6 +11,7 @@ use litellm_python_interop::{ use super::callbacks; use super::errors::to_pyerr as ocr_error_to_pyerr; use super::project::{ProjectedOcrFields, admitted_call, project_request}; +use crate::lifecycle::contract::{CallMode, CallbackPhase, PythonCallType}; use crate::lifecycle::{ OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call, }; @@ -99,7 +100,11 @@ impl PythonOcrHost { &projected.fields.secret_fields, &request.url, )?; - if !self.state.logger()?.callbacks_needed(py, "payload")? { + if !self + .state + .logger()? + .callbacks_needed(py, CallbackPhase::Payload)? + { self.state .logger()? .object(py) @@ -147,7 +152,7 @@ impl PythonOcrHost { request: OcrPostCallRequest, ) -> PyResult { let logger = self.state.logger()?; - if logger.callbacks_needed(py, "payload")? { + if logger.callbacks_needed(py, CallbackPhase::Payload)? { let projected = self.projected()?; logger.post_ocr( py, @@ -301,8 +306,12 @@ fn _ocr_lifecycle( py, args.unbind(), kwargs.copy()?.unbind(), - asynchronous, - if asynchronous { "aocr" } else { "ocr" }, + CallMode::from_async(asynchronous), + if asynchronous { + PythonCallType::AsyncOcr + } else { + PythonCallType::Ocr + }, )?, adapter: host.unbind(), data: OcrHostData::Unprojected { diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs index 713093128ba..69a7180833a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs @@ -1,10 +1,11 @@ -use litellm_core::Error; use std::future::Future; -use litellm_core::ocr::wire::{OcrWireRequest, decode_request}; use pyo3::prelude::*; use serde_json::Value; +use litellm_core::Error; +use litellm_core::ocr::wire::{OcrWireRequest, decode_request, validate_document_url}; + use super::errors::to_pyerr as ocr_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; @@ -28,6 +29,9 @@ fn prepare_ocr( .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))? .unwrap_or_default(); + validate_document_url(&document) + .map_err(Error::from) + .map_err(ocr_error_to_pyerr)?; crate::errors::admit(litellm_core::ocr::admit_value( &options.model, options.custom_llm_provider.as_deref(), diff --git a/litellm-rust/crates/python-bridge/src/routes/transcription/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/transcription/lifecycle.rs index c7610a0f952..901ab8ce444 100644 --- a/litellm-rust/crates/python-bridge/src/routes/transcription/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/transcription/lifecycle.rs @@ -1,2 +1,52 @@ -// TODO: implement transcription lifecycle checkpoints before replacing the Python lifecycle -unimplemented_lifecycle_route!(Transcription, _transcription_lifecycle); +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use litellm_core::audio_transcription::lifecycle::{ + AudioTranscriptionRoute, OwnedAudioTranscriptionRequest, +}; +use litellm_python_interop::from_py_preserving_errors as from_py; + +use crate::lifecycle::completed::{self, PythonCompletedRoute}; +use crate::lifecycle::contract::{PythonCallType, RequestField}; +use crate::lifecycle::request::{object, optional_string, options, required}; + +impl PythonCompletedRoute for AudioTranscriptionRoute { + const SYNC_CALL_TYPE: PythonCallType = PythonCallType::Transcription; + const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AsyncTranscription; + + fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> { + let audio = from_py(&required(request, RequestField::Audio)?)?; + crate::errors::admit(litellm_core::audio_transcription::admit( + &required(request, RequestField::Model)?.extract::()?, + optional_string(request, RequestField::CustomLlmProvider)?.as_deref(), + &audio, + )) + } + + fn project(request: &Bound<'_, PyDict>) -> PyResult { + Ok(OwnedAudioTranscriptionRequest { + options: options(request)?, + audio: from_py(&required(request, RequestField::Audio)?)?, + optional_params: object(request, RequestField::OptionalParams)?, + }) + } +} + +#[pyfunction] +fn _transcription_lifecycle( + py: Python<'_>, + request: Bound<'_, PyDict>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, + host: Bound<'_, PyAny>, +) -> PyResult> { + completed::run::(py, request, args, kwargs, asynchronous, host) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + crate::routes::definition::add_function( + module, + wrap_pyfunction!(_transcription_lifecycle, module)?, + ) +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index ab498fb10b0..c1b2b1bebbd 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,12 +1,53 @@ from asyncio import Future from collections.abc import Callable, Coroutine -from typing import Literal, final, overload +from typing import Literal, Protocol, TypedDict, final, overload -from typing_extensions import Never +from typing_extensions import Never, NotRequired, Required from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.lifecycle import Await, Complete from litellm.rust_bridge.ocr import LiteLLMOcrRequest +class _CommonLifecycleRequest(TypedDict): + model: Required[str] + api_key: NotRequired[str | None] + api_base: NotRequired[str | None] + custom_llm_provider: NotRequired[str | None] + extra_headers: NotRequired[object] + timeout: NotRequired[object] + +class _MessagesLifecycleRequest(_CommonLifecycleRequest): + body: Required[object] + has_agentic_hook: NotRequired[bool | None] + +class _ChatCompletionsLifecycleRequest(_CommonLifecycleRequest): + messages: Required[object] + optional_params: NotRequired[object] + host_facts: NotRequired[object] + +class _TranscriptionLifecycleRequest(_CommonLifecycleRequest): + audio: Required[object] + optional_params: NotRequired[object] + +class _CompletedLifecycleHost(Protocol): + def invoke( + self, + operation: Literal[ + "post_process", + "cache_response", + "cached_response", + "project", + "before_request", + "after_response", + "response", + "map_failure", + ], + payload: object, + request: object, + kwargs: dict[str, object], + logger: object, + ) -> Await | Complete: ... + class RustBridgeDeclined(Exception): ... class RustBridgeUnavailable(Exception): ... class RustHostCallbackError(Exception): ... @@ -28,15 +69,54 @@ def _ocr_lifecycle( asynchronous: Literal[True], host: object, ) -> Coroutine[object, object, OCRResponse]: ... +@overload def _messages_lifecycle( - request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object -) -> Never: ... + request: _MessagesLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[False], + host: _CompletedLifecycleHost, +) -> object: ... +@overload +def _messages_lifecycle( + request: _MessagesLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[True], + host: _CompletedLifecycleHost, +) -> Coroutine[object, object, object]: ... +@overload def _chat_completions_lifecycle( - request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object -) -> Never: ... + request: _ChatCompletionsLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[False], + host: _CompletedLifecycleHost, +) -> object: ... +@overload +def _chat_completions_lifecycle( + request: _ChatCompletionsLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[True], + host: _CompletedLifecycleHost, +) -> Coroutine[object, object, object]: ... +@overload def _transcription_lifecycle( - request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object -) -> Never: ... + request: _TranscriptionLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[False], + host: _CompletedLifecycleHost, +) -> object: ... +@overload +def _transcription_lifecycle( + request: _TranscriptionLifecycleRequest, + args: tuple[object, ...], + kwargs: dict[str, object], + asynchronous: Literal[True], + host: _CompletedLifecycleHost, +) -> Coroutine[object, object, object]: ... def _embeddings_lifecycle( request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object ) -> Never: ... @@ -181,11 +261,6 @@ def gil_stats() -> dict[str, int]: ... __all__ = [ "_OCR_MAX_FILE_BYTES", - "ResponsesWebSocketConnection", - "RustBridgeDeclined", - "RustBridgeUnavailable", - "RustHostCallbackError", - "RustUpstreamError", "_chat_completions_lifecycle", "_embeddings_lifecycle", "_image_edit_lifecycle", @@ -200,6 +275,11 @@ __all__ = [ "_responses_lifecycle", "_speech_lifecycle", "_transcription_lifecycle", + "ResponsesWebSocketConnection", + "RustBridgeDeclined", + "RustBridgeUnavailable", + "RustHostCallbackError", + "RustUpstreamError", "achat_completions", "amessages", "aocr", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index a3ce25ccfdb..9c883092b1e 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -113,7 +113,7 @@ COMPONENTS: Final[Mapping[ComponentName, NativeComponent]] = MappingProxyType( ), ComponentName.CHAT_COMPLETIONS: _component( ComponentName.CHAT_COMPLETIONS, - _experimental_completed, + _experimental(), ("chat_completions", "achat_completions", "_chat_completions_lifecycle"), ), ComponentName.TRANSCRIPTION: _component( diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 3d20d6a249f..c4862ba52ff 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -113,13 +113,19 @@ def test_decline_has_no_logging_effect_and_runs_one_fallback(monkeypatch: pytest assert events == [] -def test_streaming_uses_python_without_loading_native() -> None: - native_call: Final = _RecordingCall() +def test_streaming_decline_comes_from_the_native_call(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_native_bridge(monkeypatch) + native_call: Final = _RecordingCall(error=_FakeDeclined("streaming")) bridge.set_rust_chat_completions(chat_completions=native_call) kwargs: Final = _call_kwargs(ModelResponse()) kwargs["stream"] = True assert bridge.chat_completions(**kwargs) == "python" - assert native_call.calls == [] + assert len(native_call.calls) == 1 + assert native_call.calls[0]["host_facts"] == { + "stream": True, + "anthropic_user_id": False, + "bedrock_metadata_owned": False, + } def test_host_facts_reach_the_single_native_call(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_litellm_rust/test_route_foundation.py b/tests/test_litellm_rust/test_route_foundation.py index cc77bf808c1..35f4a41ceb7 100644 --- a/tests/test_litellm_rust/test_route_foundation.py +++ b/tests/test_litellm_rust/test_route_foundation.py @@ -7,26 +7,20 @@ import pytest from litellm.rust_bridge import _native from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import NATIVE_EXPORTS -from litellm.rust_bridge.chat_completions.lifecycle import LIFECYCLE as CHAT_COMPLETIONS -from litellm.rust_bridge.configuration import ExecutionDecision, ComponentName +from litellm.rust_bridge.configuration import ComponentName, ExecutionDecision from litellm.rust_bridge.embeddings.lifecycle import LIFECYCLE as EMBEDDINGS from litellm.rust_bridge.image_edit.lifecycle import LIFECYCLE as IMAGE_EDIT from litellm.rust_bridge.image_generation.lifecycle import LIFECYCLE as IMAGE_GENERATION -from litellm.rust_bridge.messages.lifecycle import LIFECYCLE as MESSAGES from litellm.rust_bridge.moderation.lifecycle import LIFECYCLE as MODERATION from litellm.rust_bridge.rerank.lifecycle import LIFECYCLE as RERANK from litellm.rust_bridge.responses.lifecycle import LIFECYCLE as RESPONSES from litellm.rust_bridge.route import ComponentExecution, NativeLifecycle from litellm.rust_bridge.runtime import BridgeErrorContext, invoke from litellm.rust_bridge.speech.lifecycle import LIFECYCLE as SPEECH -from litellm.rust_bridge.transcription.lifecycle import LIFECYCLE as TRANSCRIPTION pytestmark = pytest.mark.requires_rust_extension UNIMPLEMENTED: Final[dict[ComponentName, NativeBinding[NativeLifecycle[object, object]]]] = { - ComponentName.MESSAGES: MESSAGES, - ComponentName.CHAT_COMPLETIONS: CHAT_COMPLETIONS, - ComponentName.TRANSCRIPTION: TRANSCRIPTION, ComponentName.EMBEDDINGS: EMBEDDINGS, ComponentName.RERANK: RERANK, ComponentName.IMAGE_GENERATION: IMAGE_GENERATION, @@ -126,6 +120,16 @@ def test_transcription_declines_audio_format_before_credentials() -> None: _native.transcription("model", {"format": "unsupported", "data": "YQ=="}, custom_llm_provider="bedrock") +def test_transcription_lifecycle_declines_audio_format_before_host_work() -> None: + request: Final = { + "model": "model", + "audio": {"format": "unsupported", "data": "YQ=="}, + "custom_llm_provider": "bedrock", + } + with pytest.raises(_native.RustBridgeDeclined, match="audio format"): + _native._transcription_lifecycle(request, (), {}, False, UntouchedInput()) + + def test_websocket_declines_before_parsing_or_dialing_url() -> None: with pytest.raises(_native.RustBridgeDeclined): _native.ResponsesWebSocketConnection.connect("not a URL", custom_llm_provider="azure")